#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>

#include <chrono>

#include <iostream>

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 void GetWindowPID(xcb_connection_t* connection, const std::vector<xcb_window_t>& windows, std::unordered_map<xcb_window_t, pid_t>& result) {
	std::vector<xcb_res_query_client_ids_cookie_t> cookies;
	cookies.reserve(windows.size());

	for (const auto& window : windows) {
		xcb_res_client_id_spec_t spec = {
			.client = window,
			.mask = XCB_RES_CLIENT_ID_MASK_LOCAL_CLIENT_PID
		};

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

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

		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) {
				result[windows.at(i)] = *::xcb_res_client_id_value_value(it.data);
				continue;
			}
		}
	}
}

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

/* I have no idea why this works. */
static bool WalkWindows(xcb_connection_t* connection, const std::vector<xcb_window_t>& roots, std::set<xcb_window_t>& result) {
	/* move this somewhere */
	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;

	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 (const auto& cookie : cookies) {
		xcb_query_tree_reply_t* reply = ::xcb_query_tree_reply(connection, cookie, NULL);

		std::vector<xcb_window_t> windows = [reply]{
			xcb_window_t* windows = ::xcb_query_tree_children(reply);
			int len = ::xcb_query_tree_children_length(reply);

			std::vector<xcb_window_t> w;
			w.reserve(len);

			for (int i = 0; i < len; i++)
				w.push_back(windows[i]);

			return w;
		}();

		std::vector<xcb_get_property_cookie_t> state_property_cookies;
		state_property_cookies.reserve(windows.size());

		for (int i = 0; i < windows.size(); i++)
			state_property_cookies.push_back(::xcb_get_property(connection, 0, windows[i], Atom_WM_STATE, Atom_WM_STATE, 0, 0));

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

			if (reply->format || reply->type || reply->length) {
				result.insert(windows[i]);
				free(reply);
				continue;
			}

			free(reply);
		}

		if (WalkWindows(connection, windows, result))
			continue;
	}

	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::set<xcb_window_t> windows;
	{
		std::vector<xcb_window_t> roots;
		{
			xcb_screen_iterator_t iter = ::xcb_setup_roots_iterator(xcb_get_setup(connection));
			for (; iter.rem; ::xcb_screen_next(&iter))
				roots.push_back(iter.data->root);
		}

		
		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_get_property_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));
		pid_property_cookies.push_back(::xcb_get_property(connection, 0, window, XCB_ATOM_WM_NAME, XCB_ATOM_CARDINAL, 0L, 1L));
	}

	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) {
				const char* data = reinterpret_cast<char*>(::xcb_get_property_value(reply));
				int len = ::xcb_get_property_value_length(reply);

				win.text = std::string((char*)data, len);
			}
			free(reply);
		}
		Process proc = {0};
		{
			xcb_get_property_reply_t* reply = ::xcb_get_property_reply(connection, pid_property_cookies.at(i), NULL);
			if (reply)
				proc.pid = *reinterpret_cast<uint32_t*>(::xcb_get_property_value(reply));

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

	::xcb_disconnect(connection);

	return true;
}

}