#include "animone/win/x11.h"
#include "animone.h"
#include "animone/fd.h" /* GetProcessName() */
#include "animone/win.h"

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

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

#include <chrono>

#include <iostream>

/* This uses XCB (and it uses it *right*), so it should be plenty fast */

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

namespace animone::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());
		std::unique_ptr<xcb_intern_atom_reply_t> reply(::xcb_intern_atom_reply(connection, cookie, NULL));

		xcb_atom_t atom = reply->atom;

		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) {
		std::unique_ptr<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.get()));
			int len = ::xcb_get_property_value_length(reply.get());

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

			success |= true;
		}
	}

	return success;
}

/* This is called on every window. What this does is:
 * 1. Gets the tree of children
 * 2. Searches all children recursively for a WM_STATE property
 * 3. If that failed... return the original window
 */
static bool WalkWindows(xcb_connection_t* connection, int depth, xcb_atom_t Atom_WM_STATE, const xcb_window_t* windows,
                        int windows_len, std::set<xcb_window_t>& result) {
	/* The depth we should start returning at. */
	static constexpr int CUTOFF = 1;

	bool success = false;

	std::vector<xcb_query_tree_cookie_t> cookies;
	cookies.reserve(windows_len);

	for (int i = 0; i < windows_len; i++)
		cookies.push_back(::xcb_query_tree(connection, windows[i]));

	for (const auto& cookie : cookies) {
		std::unique_ptr<xcb_query_tree_reply_t> query_tree_reply(::xcb_query_tree_reply(connection, cookie, NULL));

		xcb_window_t* tree_children = ::xcb_query_tree_children(query_tree_reply.get());
		int tree_children_len = ::xcb_query_tree_children_length(query_tree_reply.get());

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

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

		for (int i = 0; i < tree_children_len; i++) {
			std::unique_ptr<xcb_get_property_reply_t> get_property_reply(
			    ::xcb_get_property_reply(connection, state_property_cookies[i], NULL));

			/* X11 is unfriendly here. what this means is "did the property exist?" */
			if (get_property_reply->format || get_property_reply->type || get_property_reply->length) {
				result.insert(tree_children[i]);
				if (depth >= CUTOFF)
					return true;

				success |= true;
				continue;
			}
		}

		if (WalkWindows(connection, depth + 1, Atom_WM_STATE, tree_children, tree_children_len, result)) {
			success |= true;
			if (depth >= CUTOFF)
				return true;
			continue;
		}
	}

	return success;
}

static bool GetAllTopLevelWindowsICCCM(xcb_connection_t* connection, const std::vector<xcb_window_t>& roots,
                                       std::set<xcb_window_t>& result) {
	bool success = false;

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

	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)
		success |= WalkWindows(connection, 0, Atom_WM_STATE, roots.data(), roots.size(), result);

	return success;
}

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

	xcb_connection_t* connection = ::xcb_connect(NULL, NULL);
	if (xcb_connection_has_error(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))
			GetAllTopLevelWindowsICCCM(connection, roots, windows);
	}

	struct WindowCookies {
		xcb_window_t window;
		xcb_get_property_cookie_t class_property_cookie;
		xcb_get_property_cookie_t name_property_cookie;
		xcb_res_query_client_ids_cookie_t pid_property_cookie;
	};

	std::vector<WindowCookies> window_cookies;
	window_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};

		WindowCookies window_cookie = {
		    window, ::xcb_get_property(connection, 0, window, XCB_ATOM_WM_CLASS, XCB_ATOM_STRING, 0L, 2048L),
		    ::xcb_get_property(connection, 0, window, XCB_ATOM_WM_NAME, XCB_ATOM_STRING, 0L, UINT_MAX),
		    ::xcb_res_query_client_ids(connection, 1, &spec)};

		window_cookies.push_back(window_cookie);
	}

	for (const auto& window_cookie : window_cookies) {
		Window win = {0};
		win.id = window_cookie.window;
		{
			/* Class name */
			std::unique_ptr<xcb_get_property_reply_t> reply(
			    ::xcb_get_property_reply(connection, window_cookie.class_property_cookie, NULL));

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

				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)));
			}
		}
		{
			/* Title text */
			std::unique_ptr<xcb_get_property_reply_t> reply(
			    ::xcb_get_property_reply(connection, window_cookie.name_property_cookie, NULL));

			if (reply) {
				const char* data = reinterpret_cast<const char*>(::xcb_get_property_value(reply.get()));
				int len = ::xcb_get_property_value_length(reply.get());

				win.text = std::string(data, len);
			}
		}
		Process proc = {0};
		{
			/* PID */
			std::unique_ptr<xcb_res_query_client_ids_reply_t> reply(
			    ::xcb_res_query_client_ids_reply(connection, window_cookie.pid_property_cookie, NULL));

			if (reply) {
				xcb_res_client_id_value_iterator_t it = ::xcb_res_query_client_ids_ids_iterator(reply.get());
				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);
						GetProcessName(proc.pid, proc.name); /* fill this in if we can */
						break;
					}
				}
			}
		}

		std::cout << "got window: " << win.id << "\n"
			<< "class name: " << win.class_name << "\n"
			<< "title: " << win.text << "\n"
			<< "PID: " << proc.pid << "\n"
			<< "executable: " << proc.name << "\n" << std::endl;

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

	::xcb_disconnect(connection);

	return true;
}

} // namespace animone::internal::x11