/*
 * win/x11.cc: provides support for X11 clients via XCB
 *
 * This code is fairly fast (I think...). As a result,
 * a lot of it is hard to read if you're unfamiliar with
 * asynchronous programming, but it works much better than
 * Xlib (which can take much longer).
*/
#include "animia/win/x11.h"
#include "animia/win.h"
#include "animia.h"

#include <xcb/xcb.h>
#include <xcb/res.h>

#include <cstdint>
#include <climits>
#include <cstring>
#include <string>
#include <set>

static size_t str_nlen(const char* s, size_t len) {
    size_t i = 0;
    for (; i < len && s[i]; i++);
    return i;
}

namespace animia::internal::x11 {

static bool GetAllTopLevelWindowsEWMH(xcb_connection_t* connection, const std::vector<xcb_window_t>& roots, std::set<xcb_window_t>& result) {
	const xcb_atom_t Atom__NET_CLIENT_LIST = [connection]{
		static constexpr std::string_view name = "_NET_CLIENT_LIST";
		xcb_intern_atom_cookie_t cookie = ::xcb_intern_atom(connection, true, name.size(), name.data());
		xcb_intern_atom_reply_t* reply = ::xcb_intern_atom_reply(connection, cookie, NULL);

		xcb_atom_t atom = reply->atom;
		free(reply);
		return atom;
	}();
	if (Atom__NET_CLIENT_LIST == XCB_ATOM_NONE)
		return false; // BTFO

	bool success = false;

	std::vector<xcb_get_property_cookie_t> cookies;
	cookies.reserve(roots.size());

	for (const auto& root : roots)
		cookies.push_back(::xcb_get_property(connection, 0, root, Atom__NET_CLIENT_LIST, XCB_ATOM_ANY, 0L, UINT_MAX));

	for (const auto& cookie : cookies) {
		xcb_get_property_reply_t* reply = ::xcb_get_property_reply(connection, cookie, NULL);

		if (reply && reply->length == 32) {
			xcb_window_t* value = reinterpret_cast<xcb_window_t*>(::xcb_get_property_value(reply));
			int len = ::xcb_get_property_value_length(reply);

			for (size_t i = 0; i < len; i++)
				result.insert(value[i]);

			success = true;
		}

		free(reply);
	}

	return success;
}

static bool WalkWindows(xcb_connection_t* connection, const std::vector<xcb_window_t>& roots, std::set<xcb_window_t>& result) {
	/* move this somewhere else pl0x */
	xcb_atom_t Atom_WM_STATE = [connection]{
		static constexpr std::string_view name = "WM_STATE";
		xcb_intern_atom_cookie_t cookie = ::xcb_intern_atom(connection, true, name.size(), name.data());
		xcb_intern_atom_reply_t* reply = ::xcb_intern_atom_reply(connection, cookie, NULL);

		xcb_atom_t atom = reply->atom;
		free(reply);
		return atom;
	}();
	if (Atom_WM_STATE == XCB_ATOM_NONE)
		return false;

	/* for each toplevel: search recursively for */
	std::function<bool(const xcb_window_t*, const int, xcb_window_t&)> find_wm_state_window = [&](const xcb_window_t* wins, const int wins_len, xcb_window_t& out) {
		/* Check for wm_state */
		{
			std::vector<xcb_get_property_cookie_t> property_cookies;
			property_cookies.reserve(wins_len);

			for (size_t j = 0; j < wins_len; j++)
				property_cookies.push_back(::xcb_get_property(connection, 0, wins[j], Atom_WM_STATE, Atom_WM_STATE, 0, 0));

			for (size_t j = 0; j < property_cookies.size(); j++) {
				xcb_generic_error_t* err = NULL;
				xcb_get_property_reply_t* reply = ::xcb_get_property_reply(connection, property_cookies.at(j), &err);

				if (reply->format || reply->type || reply->length) {
					out = wins[j];
					free(reply);
					return true;
				}

				free(reply);
			}
		}

		/* Query tree for recursion */
		{
			std::vector<xcb_query_tree_cookie_t> cookies;
			cookies.reserve(wins_len);

			for (size_t j = 0; j < wins_len; j++)
				cookies.push_back(::xcb_query_tree(connection, wins[j]));

			for (const auto& cookie : cookies) {
				xcb_query_tree_reply_t* reply = ::xcb_query_tree_reply(connection, cookie, NULL);

				xcb_window_t* windows = ::xcb_query_tree_children(reply);
				int len = ::xcb_query_tree_children_length(reply);

				if (find_wm_state_window(windows, len, out)) {
					free(reply);
					return true;
				}

				free(reply);
			}
		}

		return false;
	};

	/* Get the tree for each root */
	std::vector<xcb_query_tree_cookie_t> cookies;
	cookies.reserve(roots.size());

	for (const auto& root : roots)
		cookies.push_back(::xcb_query_tree(connection, root));

	for (size_t i = 0; i < cookies.size(); i++) {
		xcb_query_tree_reply_t* reply = ::xcb_query_tree_reply(connection, cookies.at(i), NULL);

		xcb_window_t* windows = ::xcb_query_tree_children(reply);
		int len = ::xcb_query_tree_children_length(reply);
		if (len < 1)
			continue;

		/* Then get the tree of each child window. */
		std::vector<xcb_query_tree_cookie_t> cookies;
		cookies.reserve(len);

		for (size_t j = 0; j < len; j++)
			cookies.push_back(::xcb_query_tree(connection, windows[j]));

		/* For each child window... */
		for (size_t j = 0; j < cookies.size(); j++) {
			xcb_query_tree_reply_t* reply = ::xcb_query_tree_reply(connection, cookies.at(j), NULL);

			xcb_window_t* children = ::xcb_query_tree_children(reply);
			int children_len = ::xcb_query_tree_children_length(reply);
			if (children_len < 1) {
				result.insert(windows[j]);
				continue;
			}

			xcb_window_t out = windows[j];
			/* Search recursively for a window with WM_STATE. If we don't,
			 * just add the toplevel.
			*/
			find_wm_state_window(children, children_len, out);
			result.insert(out);
		}

		free(reply);
	}

	return false;
}

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

	xcb_connection_t* connection = ::xcb_connect(NULL, NULL);
	if (!connection)
		return false;

	std::vector<xcb_screen_t> screens;

	{
		xcb_screen_iterator_t iter = ::xcb_setup_roots_iterator(xcb_get_setup(connection));
		for (; iter.rem; ::xcb_screen_next(&iter))
			screens.push_back(*iter.data);
	}

	std::vector<xcb_window_t> roots;
	roots.reserve(screens.size());

	for (const auto& screen : screens)
		roots.push_back(screen.root);

	std::set<xcb_window_t> windows;
	if (!GetAllTopLevelWindowsEWMH(connection, roots, windows))
		WalkWindows(connection, roots, windows);

	std::vector<xcb_get_property_cookie_t> class_property_cookies;
	std::vector<xcb_get_property_cookie_t> name_property_cookies;
	std::vector<xcb_res_query_client_ids_cookie_t> pid_property_cookies;
	class_property_cookies.reserve(windows.size());
	name_property_cookies.reserve(windows.size());
	pid_property_cookies.reserve(windows.size());

	for (const auto& window : windows) {
		class_property_cookies.push_back(::xcb_get_property(connection, 0, window, XCB_ATOM_WM_CLASS, XCB_ATOM_STRING, 0L, 2048L));
		name_property_cookies.push_back(::xcb_get_property(connection, 0, window, XCB_ATOM_WM_NAME, XCB_ATOM_STRING, 0L, UINT_MAX));

		xcb_res_client_id_spec_t spec = {
			.client = window,
			.mask = XCB_RES_CLIENT_ID_MASK_LOCAL_CLIENT_PID
		};

		pid_property_cookies.push_back(::xcb_res_query_client_ids(connection, 1, &spec));
	}

	size_t i = 0;
	for (const auto& window : windows) {
		Window win = {0};
		win.id = window;
		{
			xcb_get_property_reply_t* reply = ::xcb_get_property_reply(connection, class_property_cookies.at(i), NULL);

			if (reply && reply->format == 8) {
				const char* data = reinterpret_cast<char*>(::xcb_get_property_value(reply));
				const int data_len = ::xcb_get_property_value_length(reply);

				int instance_len = str_nlen(data, data_len);
				const char* class_name = data + instance_len + 1;

				win.class_name = std::string(class_name, str_nlen(class_name, data_len - (instance_len + 1)));
			}

			free(reply);
		}
		{
			xcb_get_property_reply_t* reply = ::xcb_get_property_reply(connection, name_property_cookies.at(i), NULL);

			if (reply && reply->format == 8) {
				const char* data = reinterpret_cast<char*>(::xcb_get_property_value(reply));
				int len = ::xcb_get_property_value_length(reply);

				win.text = std::string(data, len);
			}

			free(reply);
		}
		Process proc = {0};
		{
			xcb_res_query_client_ids_reply_t* reply = ::xcb_res_query_client_ids_reply(connection, pid_property_cookies.at(i), NULL);

			if (reply->length) {
				xcb_res_client_id_value_iterator_t it = ::xcb_res_query_client_ids_ids_iterator(reply);
				for (; it.rem; ::xcb_res_client_id_value_next(&it)) {
					if (it.data->spec.mask & XCB_RES_CLIENT_ID_MASK_LOCAL_CLIENT_PID) {
						proc.pid = *::xcb_res_client_id_value_value(it.data);
						break;
					}
				}
			}

			free(reply);
		}

		if (!window_proc(proc, win)) {
			::xcb_disconnect(connection);
			return false;
		}
		i++;
	}

	::xcb_disconnect(connection);

	return true;
}

}