#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_*
#ifdef HAVE_XRES
#include <X11/extensions/XRes.h>
#endif

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

namespace animia::internal::x11 {

/* specify that these are X types. */
typedef ::Window XWindow;
typedef ::Display XDisplay;
typedef ::Atom XAtom;

/* should return UTF8_STRING or STRING. this means we are not
 * *guaranteed* a UTF-8 string back.
*/
static bool GetWindowPropertyAsString(XDisplay* display, XWindow window, XAtom atom, std::string& result) {
	if (atom == None)
		return false;

	XAtom Atom_UTF8_STRING = ::XInternAtom(display, "UTF8_STRING", False);

	int format;
	unsigned long leftover_bytes, num_of_items;
	XAtom type;
	unsigned char* data;

	int status = ::XGetWindowProperty(display, window, atom, 0L, (~0L), False, AnyPropertyType,
	                                  &type, &format, &num_of_items, &leftover_bytes, &data);
	if (status != Success || !(type == Atom_UTF8_STRING || type == XA_STRING) || !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(XDisplay* display, XWindow window, pid_t& result) {
#ifdef HAVE_XRES
	{
		long num_ids;
		XResClientIdValue *client_ids;
		XResClientIdSpec spec = {
			.client = window,
			.mask = XRES_CLIENT_ID_PID_MASK
		};

		::XResQueryClientIds(display, 1, &spec, &num_ids, &client_ids);

		for (long i = 0; i < num_ids; i++) {
			if (client_ids[i].spec.mask == XRES_CLIENT_ID_PID_MASK) {
				result = ::XResGetClientPid(&client_ids[i]);
				::XResClientIdsDestroy(num_ids, client_ids);
				return true;
			}
		}

		::XResClientIdsDestroy(num_ids, client_ids);

		return false;
	}
#endif

	XAtom Atom__NET_WM_PID = ::XInternAtom(display, "_NET_WM_PID", True);
	if (Atom__NET_WM_PID == None)
		return false;

	int format;
	unsigned long leftover_bytes, num_of_items;
	XAtom type;
	unsigned char* data;

	int status = ::XGetWindowProperty(display, window, Atom__NET_WM_PID, 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>(*reinterpret_cast<uint32_t*>(data));

	::XFree(data);

	return true;
}

static bool FetchName(XDisplay* display, XWindow window, std::string& result) {
	if (GetWindowPropertyAsString(display, window, ::XInternAtom(display, "_NET_WM_NAME", True), result))
		return true;

	if (GetWindowPropertyAsString(display, window, ::XInternAtom(display, "WM_NAME", True), result))
		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 GetAllTopLevelWindowsEWMH(XDisplay* display, XWindow root, std::set<XWindow>& result) {
	XAtom Atom__NET_CLIENT_LIST = XInternAtom(display, "_NET_CLIENT_LIST", True);
	if (Atom__NET_CLIENT_LIST == None)
		return false;

	XAtom actual_type;
	int format;
	unsigned long num_of_items, bytes_after;
	unsigned char* data = nullptr;

	{
		int status = ::XGetWindowProperty(
			display, root, Atom__NET_CLIENT_LIST,
			0L, (~0L), false, AnyPropertyType, &actual_type,
			&format, &num_of_items, &bytes_after, &data
		);

		if (status < Success || !num_of_items)
			return false;
	}

	XWindow* arr = (XWindow*)data;

	for (uint32_t i = 0; i < num_of_items; i++)
		result.insert(arr[i]);

	::XFree(data);

	return true;
}

static bool GetAllTopLevelWindows(XDisplay* display, XWindow root, std::set<XWindow>& result) {
	// EWMH. Takes about 15 ms on a fairly good PC.
	if (GetAllTopLevelWindowsEWMH(display, root, result))
		return true;

	// Fallback to ICCCM. Takes about the same time on a good PC.
	XAtom Atom_WM_STATE = XInternAtom(display, "WM_STATE", True);
	if (Atom_WM_STATE == None)
		return false;

	auto window_has_wm_state = [&](XWindow window) -> bool {
		int format;
		Atom actual_type;
		unsigned long num_of_items, bytes_after;
		unsigned char* data = nullptr;

		int status = ::XGetWindowProperty(
			display, window, Atom_WM_STATE,
			0L, (~0L), false, AnyPropertyType, &actual_type,
			&format, &num_of_items, &bytes_after, &data
		);

		::XFree(data);

		return !(actual_type == None && !format && !bytes_after);
	};

	std::function<bool(XWindow, XWindow&)> immediate_child_get_toplevel = [&](XWindow window, XWindow& result) {
		result = window;
		if (window_has_wm_state(window))
			return true;

		unsigned int num_children = 0;
		XWindow* children_arr = nullptr;

		XWindow root_return;
		XWindow parent_return;

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

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

		for (unsigned int i = 0; i < num_children; i++) {
			if (immediate_child_get_toplevel(children_arr[i], result)) {
				::XFree(children_arr);
				return true;
			}
		}

		::XFree(children_arr);
		return false;
	};

	unsigned int num_children = 0;
	XWindow* children_arr = nullptr;

	XWindow root_return;
	XWindow parent_return;

	int status = ::XQueryTree(display, root, &root_return, &parent_return, &children_arr, &num_children);
	if (!status || !children_arr)
		return false; // how

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

	for (unsigned int i = 0; i < num_children; i++) {
		XWindow res;
		if (immediate_child_get_toplevel(children_arr[i], res))
			result.insert(res);
	}

	::XFree(children_arr);

	return true;
}

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

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

	XWindow root = ::XDefaultRootWindow(display);

	std::set<XWindow> windows;
	GetAllTopLevelWindows(display, root, windows);

	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;
}

}