#include "animia/win/x11.h"
#include "animia/win.h"
#include "animia.h"

#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/Xatom.h> // XA_*

#include <cstdint>
#include <string>
#include <set>

/* The code for this is very fugly because X11 uses lots of generic type names
   (i.e., Window, Display), so I have to use :: when defining vars to distinguish
   between Animia's types and X11's types */

namespace animia::internal::x11 {

/* should return UTF8_STRING or STRING */
static bool GetWindowPropertyAsString(::Display* display, ::Window window, ::Atom atom, std::string& result, ::Atom reqtype = AnyPropertyType) {
	int format;
	unsigned long leftover_bytes, num_of_items;
	::Atom type;
	unsigned char* data;

	int status = ::XGetWindowProperty(display, window, atom, 0L, (~0L), False, reqtype,
	                                  &type, &format, &num_of_items, &leftover_bytes, &data);
	if (status != Success || !(reqtype == AnyPropertyType || type == reqtype) || !num_of_items)
		return false;

	result = std::string((char*)data, num_of_items);

	::XFree(data);

	return true;
}

/* this should return CARDINAL, a 32-bit integer */
static bool GetWindowPID(::Display* display, ::Window window, pid_t& result) {
	int format;
	unsigned long leftover_bytes, num_of_items;
	::Atom atom = ::XInternAtom(display, "_NET_WM_PID", False), type;
	unsigned char* data;

	int status = ::XGetWindowProperty(display, window, atom, 0L, (~0L), False, XA_CARDINAL,
	                                  &type, &format, &num_of_items, &leftover_bytes, &data);
	if (status != Success || type != XA_CARDINAL || num_of_items < 1)
		return false;

	result = static_cast<pid_t>(*(uint32_t*)data);

	::XFree(data);

	return true;
}

static bool FetchName(::Display* display, ::Window window, std::string& result) {
	/* TODO: Check if XInternAtom created None or not... */
	if (GetWindowPropertyAsString(display, window, ::XInternAtom(display, "_NET_WM_NAME", False),
		                          result, ::XInternAtom(display, "UTF8_STRING", False)))
		return true;

	if (GetWindowPropertyAsString(display, window, ::XInternAtom(display, "WM_NAME", False),
		                           result, XA_STRING))
		return true;

	/* Fallback to XGetWMName() */
	XTextProperty text;

	{
		int status = ::XGetWMName(display, window, &text);
		if (!status || !text.value || !text.nitems)
			return false;
	}

	char** list;

	{
		int count;

		int status = ::XmbTextPropertyToTextList(display, &text, &list, &count);
		if (status != Success || !count || !*list)
			return false;
	}

	::XFree(text.value);

	result = *list;

	::XFreeStringList(list);

	return true;
}

static bool WalkWindows(::Display* display, std::set<::Window>& children, const std::set<::Window>& windows) {
	if (windows.empty())
		return false;

	for (const ::Window& window : windows) {
		unsigned int num_children = 0;
		::Window* children_arr = nullptr;

		::Window root_return;
		::Window parent_return;

		int status = ::XQueryTree(display, window, &root_return, &parent_return, &children_arr, &num_children);
		if (!status || !children_arr)
			continue;

		if (num_children < 1) {
			::XFree(children_arr);
			continue;
		}

		for (int i = 0; i < num_children; i++)
			if (!children.count(children_arr[i]))
				children.insert(children_arr[i]);

		::XFree(children_arr);

		std::set<::Window> children_children;

		if (WalkWindows(display, children_children, children))
			children.insert(children_children.begin(), children_children.end());
	}

	return true;
}

bool X11WinTools::EnumerateWindows(window_proc_t window_proc) {
	if (!window_proc)
		return false;

	::Display* display = ::XOpenDisplay(nullptr);
	if (!display)
		return false;

	::Window root = DefaultRootWindow(display);

	std::set<::Window> windows;
	WalkWindows(display, windows, {root});

	for (const auto& window : windows) {
		Window win;
		win.id = window;
		{
			::XClassHint* hint = ::XAllocClassHint();
			if (::XGetClassHint(display, window, hint)) {
				win.class_name = hint->res_class;
				::XFree(hint);
			}
		}
		FetchName(display, window, win.text);

		Process proc;
		GetWindowPID(display, window, proc.pid);

		if (!window_proc(proc, win))
			return false;
	}

	return true;
}

}