/*
 * win/x11.c: support for X11 using XCB
 *
 * some things might be wrong here due to
 * having to use recursive logic, but whatever
 */
#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 <cassert>
#include <cstdint>
#include <cstring>
#include <cstdlib>
#include <set>
#include <string>
#include <memory>

#include <chrono>
#include <unordered_map>

#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 animone::internal::x11 {

template<typename T>
struct XcbDestructor {
	using pointer = T*;
	void operator()(pointer t) const { std::free(t); };
};

template<typename T>
using XcbPtr = std::unique_ptr<T, XcbDestructor<T>>;

/* --------------------------------------------------------------
 * atom cruft */

enum class NeededAtom {
	/* EWMH */
	NET_CLIENT_LIST,
	NET_WM_NAME,
	UTF8_STRING,

	/* ICCCM */
	WM_STATE,
};

static const std::unordered_map<NeededAtom, std::string> atom_strings = {
	{NeededAtom::NET_CLIENT_LIST, "_NET_CLIENT_LIST"},
	{NeededAtom::NET_WM_NAME, "_NET_WM_NAME"},
	{NeededAtom::UTF8_STRING, "UTF8_STRING"},

	{NeededAtom::WM_STATE, "WM_STATE"},
};

using XcbAtoms = std::unordered_map<NeededAtom, xcb_atom_t>;

static bool GetAllTopLevelWindowsEWMH(xcb_connection_t* connection, const XcbAtoms& atoms, const std::vector<xcb_window_t>& roots,
                                      std::set<xcb_window_t>& result) {
	if (atoms.at(NeededAtom::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, atoms.at(NeededAtom::NET_CLIENT_LIST), XCB_ATOM_ANY, 0L, UINT_MAX));

	for (const auto& cookie : cookies) {
		XcbPtr<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()) / sizeof(xcb_window_t);

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

			success |= true;
		}
	}

	return success;
}

/* This should be called with a list of toplevels for each root. */
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 level of depth we want to cut off past; since we want to go over each top level window,
	 * we cut off after we've passed the root window and the toplevel. */
	static constexpr int CUTOFF = 2;
	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 (int i = 0; i < cookies.size(); i++) {
		/* XXX is it *really* okay to ask xcb for a cookie and then never ask for a reply?
		 * valgrind doesn't complain, so I'm not gonna care for now. */
		XcbPtr<xcb_query_tree_reply_t> query_tree_reply(::xcb_query_tree_reply(connection, cookies[i], 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());

		/* search for any window with a WM_STATE property */
		std::vector<xcb_get_property_cookie_t> state_cookies;
		state_cookies.reserve(tree_children_len);

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

		bool found = false;

		for (int i = 0; i < tree_children_len; i++) {
			XcbPtr<xcb_get_property_reply_t> get_property_reply(::xcb_get_property_reply(connection, state_cookies[i], NULL));
			if (!get_property_reply)
				continue;

			/* did we get valid data? */
			if (get_property_reply->type == Atom_WM_STATE || get_property_reply->format != 0 || get_property_reply->bytes_after != 0) {
				int len = ::xcb_get_property_value_length(get_property_reply.get());
				if (len < sizeof(uint32_t))
					continue;

				uint32_t state = *reinterpret_cast<uint32_t*>(::xcb_get_property_value(get_property_reply.get()));
				if (state != 1) // NormalState
					continue;
				
				result.insert(tree_children[i]);
				found = true;
				if (depth >= CUTOFF)
					return true;
			}
		}

		if (found)
			continue;

		bool res = WalkWindows(connection, depth + 1, Atom_WM_STATE, tree_children, tree_children_len, result);

		if (depth >= CUTOFF)
			return res;
	}

	return true;
}

static bool GetAllTopLevelWindowsICCCM(xcb_connection_t* connection, const XcbAtoms& atoms, const std::vector<xcb_window_t>& roots,
                                       std::set<xcb_window_t>& result) {
	if (atoms.at(NeededAtom::WM_STATE) == XCB_ATOM_NONE)
		return false;

	return WalkWindows(connection, 0, atoms.at(NeededAtom::WM_STATE), roots.data(), roots.size(), result);
}

static XcbAtoms InitializeAtoms(xcb_connection_t* connection) {
	XcbAtoms atoms;

	std::unordered_map<NeededAtom, xcb_intern_atom_cookie_t> atom_cookies;

	for (const auto& [atom, str] : atom_strings)
		atom_cookies[atom] = ::xcb_intern_atom(connection, 1, str.size(), str.data());

	for (const auto& [atom, cookie] : atom_cookies) {
		XcbPtr<xcb_intern_atom_reply_t> reply(::xcb_intern_atom_reply(connection, cookie, NULL));
		if (!reply)
			atoms[atom] = XCB_ATOM_NONE;

		atoms[atom] = reply->atom;
	}

	return atoms;
}

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;

	XcbAtoms atoms = InitializeAtoms(connection);

	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, atoms, roots, windows))
			GetAllTopLevelWindowsICCCM(connection, atoms, roots, windows);
	}

	struct WindowCookies {
		xcb_window_t window;
		xcb_get_property_cookie_t class_name;
		xcb_get_property_cookie_t name_utf8;
		xcb_get_property_cookie_t name;
		xcb_res_query_client_ids_cookie_t pid;
	};

	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 = window,
		    .class_name = ::xcb_get_property(connection, 0, window, XCB_ATOM_WM_CLASS, XCB_ATOM_STRING, 0L, 2048L),
		    .name_utf8 = ::xcb_get_property(connection, 0, window, atoms[NeededAtom::NET_WM_NAME], atoms[NeededAtom::UTF8_STRING], 0L, UINT_MAX),
		    .name = ::xcb_get_property(connection, 0, window, XCB_ATOM_WM_NAME, XCB_ATOM_STRING, 0L, UINT_MAX),
		    .pid = ::xcb_res_query_client_ids(connection, 1, &spec),
		};

		window_cookies.push_back(window_cookie);
	}

	for (const auto& window_cookie : window_cookies) {
		Window win;
		win.id = window_cookie.window;
		{
			/* Class name */
			XcbPtr<xcb_get_property_reply_t> reply(::xcb_get_property_reply(connection, window_cookie.class_name, 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 */
			XcbPtr<xcb_get_property_reply_t> reply_utf8(::xcb_get_property_reply(connection, window_cookie.name_utf8, NULL));
			XcbPtr<xcb_get_property_reply_t> reply(::xcb_get_property_reply(connection, window_cookie.name, NULL));
			int utf8_len = ::xcb_get_property_value_length(reply_utf8.get());
			int len = ::xcb_get_property_value_length(reply.get());

			if (reply_utf8 && utf8_len > 0) {
				const char* data = reinterpret_cast<const char*>(::xcb_get_property_value(reply_utf8.get()));

				win.text = std::string(data, utf8_len);
			} else if (reply && len > 0) {
				const char* data = reinterpret_cast<const char*>(::xcb_get_property_value(reply.get()));

				win.text = std::string(data, len);
			}
		}
		Process proc;
		proc.platform = ExecutablePlatform::Posix; // not entirely correct, but whatever. who cares
		{
			/* PID */
			XcbPtr<xcb_res_query_client_ids_reply_t> reply(::xcb_res_query_client_ids_reply(connection, window_cookie.pid, 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.comm); /* fill this in if we can */
						break;
					}
				}
			}
		}

		/* debug printing
		std::cout << "window found: " << std::hex << win.id << std::dec << "\n"
			<< " name: " << win.text << "\n"
			<< " class: " << win.class_name << "\n"
			<< " pid: " << proc.pid << "\n"
			<< " comm: " << proc.name << std::endl;
		*/

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

	::xcb_disconnect(connection);

	return true;
}

} // namespace animone::internal::x11