#include "mkb_hook.hpp"

#include <atomic>

#include <windows.h>

#include <safetyhook.hpp>

#include "hook_registry.hpp"

namespace coop::hook
{

namespace
{

// Synthesized input state the polling hooks report. Written by the worker thread
// (mkb_pump), read by the game's thread inside the detours -> all atomic.
std::atomic<bool> g_active{false};
std::atomic<bool> g_key_down[256];        // by Win32 virtual-key (incl. VK_LBUTTON etc.)
std::atomic<long> g_cursor_x{0};          // last forwarded mouse position (game client px)
std::atomic<long> g_cursor_y{0};
std::atomic<bool> g_have_cursor{false};   // a mouse event has been forwarded at least once
std::atomic<void*> g_target{nullptr};     // game main window (HWND), resolved lazily

safetyhook::InlineHook g_hk_async;
safetyhook::InlineHook g_hk_kbstate;
safetyhook::InlineHook g_hk_cursor;

int g_id_pump = -1;
int g_id_async = -1;
int g_id_kbstate = -1;
int g_id_cursor = -1;
bool g_installed = false;

// --- Synthesized-state detours (let polling games see forwarded input) ------

SHORT WINAPI hk_GetAsyncKeyState(int vkey)
{
	const SHORT orig = g_hk_async.stdcall<SHORT>(vkey);
	if (g_active.load(std::memory_order_relaxed) && vkey >= 0 && vkey < 256 &&
	    g_key_down[vkey].load(std::memory_order_relaxed))
	{
		return static_cast<SHORT>(0x8000) | (orig & 0x1);
	}
	return orig;
}

BOOL WINAPI hk_GetKeyboardState(PBYTE state)
{
	const BOOL r = g_hk_kbstate.stdcall<BOOL>(state);
	if (r && state != nullptr && g_active.load(std::memory_order_relaxed))
	{
		for (int vk = 0; vk < 256; ++vk)
		{
			if (g_key_down[vk].load(std::memory_order_relaxed))
			{
				state[vk] |= 0x80;
			}
		}
	}
	return r;
}

BOOL WINAPI hk_GetCursorPos(LPPOINT pt)
{
	const BOOL r = g_hk_cursor.stdcall<BOOL>(pt);
	if (g_active.load(std::memory_order_relaxed) && g_have_cursor.load(std::memory_order_relaxed) && pt != nullptr)
	{
		auto* hwnd = static_cast<HWND>(g_target.load(std::memory_order_relaxed));
		if (hwnd != nullptr)
		{
			POINT c{g_cursor_x.load(std::memory_order_relaxed), g_cursor_y.load(std::memory_order_relaxed)};
			ClientToScreen(hwnd, &c); // synth state is game-client; GetCursorPos is screen-space
			*pt = c;
			return TRUE;
		}
	}
	return r;
}

// --- Target window + message synthesis --------------------------------------

struct FindCtx
{
	DWORD pid;
	HWND best;
	long area;
};

BOOL CALLBACK find_main_proc(HWND hwnd, LPARAM lparam)
{
	auto* c = reinterpret_cast<FindCtx*>(lparam);
	DWORD pid = 0;
	GetWindowThreadProcessId(hwnd, &pid);
	if (pid != c->pid || !IsWindowVisible(hwnd) || GetWindow(hwnd, GW_OWNER) != nullptr)
	{
		return TRUE;
	}
	RECT r{};
	GetWindowRect(hwnd, &r);
	const long area = (r.right - r.left) * (r.bottom - r.top);
	if (area > c->area)
	{
		c->area = area;
		c->best = hwnd;
	}
	return TRUE;
}

HWND find_main_window()
{
	FindCtx c{GetCurrentProcessId(), nullptr, 0};
	EnumWindows(&find_main_proc, reinterpret_cast<LPARAM>(&c));
	return c.best;
}

LPARAM key_lparam(UINT vk, bool key_up)
{
	const UINT scan = MapVirtualKeyW(vk, MAPVK_VK_TO_VSC);
	LPARAM lp = 1 | (static_cast<LPARAM>(scan) << 16); // repeat count 1 + scan code
	if (key_up)
	{
		lp |= (LPARAM{1} << 30) | (LPARAM{1} << 31); // previous-down + transition (key released)
	}
	return lp;
}

void set_key(UINT vk, bool down)
{
	if (vk < 256)
	{
		g_key_down[vk].store(down, std::memory_order_relaxed);
	}
}

WPARAM mouse_button_wparam()
{
	WPARAM w = 0;
	if (g_key_down[VK_LBUTTON].load(std::memory_order_relaxed))
	{
		w |= MK_LBUTTON;
	}
	if (g_key_down[VK_RBUTTON].load(std::memory_order_relaxed))
	{
		w |= MK_RBUTTON;
	}
	if (g_key_down[VK_MBUTTON].load(std::memory_order_relaxed))
	{
		w |= MK_MBUTTON;
	}
	return w;
}

void handle_mouse(const MkbEvent& ev, bool down, HWND hwnd)
{
	g_cursor_x.store(ev.x, std::memory_order_relaxed);
	g_cursor_y.store(ev.y, std::memory_order_relaxed);
	g_have_cursor.store(true, std::memory_order_relaxed);

	const UINT vk = ev.code == 0 ? VK_LBUTTON : ev.code == 1 ? VK_RBUTTON : VK_MBUTTON;
	set_key(vk, down);
	if (hwnd == nullptr)
	{
		return;
	}
	UINT msg;
	if (ev.code == 0)
	{
		msg = down ? WM_LBUTTONDOWN : WM_LBUTTONUP;
	}
	else if (ev.code == 1)
	{
		msg = down ? WM_RBUTTONDOWN : WM_RBUTTONUP;
	}
	else
	{
		msg = down ? WM_MBUTTONDOWN : WM_MBUTTONUP;
	}
	PostMessageW(hwnd, msg, mouse_button_wparam(), MAKELPARAM(ev.x, ev.y));
}

void handle_wheel(const MkbEvent& ev, HWND hwnd)
{
	g_cursor_x.store(ev.x, std::memory_order_relaxed);
	g_cursor_y.store(ev.y, std::memory_order_relaxed);
	g_have_cursor.store(true, std::memory_order_relaxed);
	if (hwnd == nullptr)
	{
		return;
	}
	POINT pt{ev.x, ev.y};
	ClientToScreen(hwnd, &pt); // WM_MOUSEWHEEL lParam is screen coordinates
	const short delta = static_cast<short>(static_cast<std::int32_t>(ev.code));
	PostMessageW(hwnd, WM_MOUSEWHEEL, MAKEWPARAM(0, delta), MAKELPARAM(pt.x, pt.y));
}

void install_user32_hook(HMODULE user32, const char* name, void* detour, safetyhook::InlineHook& slot, int id)
{
	if (user32 == nullptr)
	{
		return;
	}
	if (void* target = reinterpret_cast<void*>(GetProcAddress(user32, name)))
	{
		slot = safetyhook::create_inline(target, detour);
		if (slot)
		{
			hook_set_installed(id, true);
		}
	}
}

} // namespace

bool install_mkb_hooks(IpcClient& ipc)
{
	if (g_installed)
	{
		return true;
	}

	if (g_id_pump < 0)
	{
		g_id_pump = hook_register("MKB pump (PostMessage)", HookSubsys_Mkb);
		g_id_async = hook_register("GetAsyncKeyState", HookSubsys_Mkb);
		g_id_kbstate = hook_register("GetKeyboardState", HookSubsys_Mkb);
		g_id_cursor = hook_register("GetCursorPos", HookSubsys_Mkb);
	}

	// Fresh synthesized state so a previous session leaves no stuck keys.
	for (int vk = 0; vk < 256; ++vk)
	{
		g_key_down[vk].store(false, std::memory_order_relaxed);
	}
	g_have_cursor.store(false, std::memory_order_relaxed);
	g_target.store(nullptr, std::memory_order_relaxed);

	HMODULE user32 = GetModuleHandleW(L"user32.dll");
	install_user32_hook(user32, "GetAsyncKeyState", reinterpret_cast<void*>(&hk_GetAsyncKeyState), g_hk_async,
	                    g_id_async);
	install_user32_hook(user32, "GetKeyboardState", reinterpret_cast<void*>(&hk_GetKeyboardState), g_hk_kbstate,
	                    g_id_kbstate);
	install_user32_hook(user32, "GetCursorPos", reinterpret_cast<void*>(&hk_GetCursorPos), g_hk_cursor, g_id_cursor);

	g_active.store(true, std::memory_order_release);
	hook_set_installed(g_id_pump, true);
	g_installed = true;
	(void)ipc;
	return true;
}

void remove_mkb_hooks()
{
	if (!g_installed)
	{
		return;
	}
	g_active.store(false, std::memory_order_release);
	g_hk_async = {};   // InlineHook destructor restores the original bytes
	g_hk_kbstate = {};
	g_hk_cursor = {};
	for (int vk = 0; vk < 256; ++vk)
	{
		g_key_down[vk].store(false, std::memory_order_relaxed); // no stuck keys
	}
	g_have_cursor.store(false, std::memory_order_relaxed);
	hook_set_installed(g_id_pump, false);
	hook_set_installed(g_id_async, false);
	hook_set_installed(g_id_kbstate, false);
	hook_set_installed(g_id_cursor, false);
	g_installed = false;
}

void mkb_pump(IpcClient& ipc)
{
	MkbRing* ring = ipc.mkb_ring();
	if (ring == nullptr || !g_active.load(std::memory_order_relaxed))
	{
		return;
	}

	HWND hwnd = static_cast<HWND>(g_target.load(std::memory_order_relaxed));
	if (hwnd == nullptr || !IsWindow(hwnd))
	{
		hwnd = find_main_window();
		g_target.store(hwnd, std::memory_order_relaxed);
	}

	MkbEvent ev{};
	while (pop_mkb_event(*ring, ev))
	{
		hook_note_call(g_id_pump);
		switch (ev.type)
		{
		case Mkb_KeyDown:
			set_key(ev.code, true);
			if (hwnd != nullptr)
			{
				PostMessageW(hwnd, WM_KEYDOWN, ev.code, key_lparam(ev.code, false));
			}
			break;
		case Mkb_KeyUp:
			set_key(ev.code, false);
			if (hwnd != nullptr)
			{
				PostMessageW(hwnd, WM_KEYUP, ev.code, key_lparam(ev.code, true));
			}
			break;
		case Mkb_Char:
			if (hwnd != nullptr)
			{
				PostMessageW(hwnd, WM_CHAR, ev.code, 1);
			}
			break;
		case Mkb_MouseDown:
			handle_mouse(ev, true, hwnd);
			break;
		case Mkb_MouseUp:
			handle_mouse(ev, false, hwnd);
			break;
		case Mkb_Wheel:
			handle_wheel(ev, hwnd);
			break;
		default:
			break;
		}
	}
}

} // namespace coop::hook