#include "mkb_hook.hpp"

#include <atomic>

#include <windows.h>

#define DIRECTINPUT_VERSION 0x0800
#include <dinput.h>

#include <safetyhook.hpp>

#include "hook_guard.hpp"
#include "hook_registry.hpp"
#include "vtable_hook.hpp"

namespace coop::hook
{

namespace
{

DetourGate g_gate; // drains in-flight polling detours before remove tears the hooks down

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

// --- DirectInput forwarding -------------------------------------------------
// Many games read keyboard/mouse via IDirectInputDevice8::GetDeviceState instead of the polling
// APIs above. All DI devices share one vtable (same coclass), so a single vtable swap -- read from
// a kept-alive probe device -- intercepts the game's existing devices too (it may have created them
// before we injected). We must use a vtable swap, not an inline hook (the x86 COM-prologue trap).
VtableHook g_vh_di_getstate;
IDirectInput8W* g_di_probe = nullptr;
IDirectInputDevice8W* g_di_probe_kbd = nullptr;
int g_id_di_getstate = -1;
// IDirectInputDevice8 vtable: IUnknown 0-2, GetCapabilities 3, EnumObjects 4, GetProperty 5,
// SetProperty 6, Acquire 7, Unacquire 8, GetDeviceState 9, GetDeviceData 10.
constexpr unsigned kIdx_IDirectInputDevice8_GetDeviceState = 9;
std::atomic<long> g_di_mouse_last_x{0}; // last forwarded cursor, for relative DI mouse deltas
std::atomic<long> g_di_mouse_last_y{0};
std::atomic<bool> g_di_mouse_primed{false};

// --- Raw Input forwarding ---------------------------------------------------
// Games that read via Raw Input (RegisterRawInputDevices + WM_INPUT -> GetRawInputData) get no
// WM_INPUT while unfocused, so we synthesize it: mkb_pump posts WM_INPUT with lParam = the address
// of a synthetic RAWINPUT slot, and this hook serves that slot's data when the game reads it back.
safetyhook::InlineHook g_hk_getrawinputdata; // user32!GetRawInputData
int g_id_rawinput = -1;
constexpr int kRawSlots = 64; // small ring of synthetic events (games consume WM_INPUT promptly)
RAWINPUT g_raw_slots[kRawSlots] = {};
std::atomic<unsigned> g_raw_head{0};

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

SHORT WINAPI hk_GetAsyncKeyState(int vkey)
{
	DetourGate::Guard guard(g_gate);
	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)
{
	DetourGate::Guard guard(g_gate);
	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)
{
	DetourGate::Guard guard(g_gate);
	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;
}

// --- DirectInput GetDeviceState detour --------------------------------------

// Win32 virtual-key -> DirectInput key index (DIK_*, a set-1 scan code). Good for the common keys
// (letters/digits/WASD/space/enter); a few extended keys (arrows) differ by the extended bit.
BYTE vk_to_dik(int vk)
{
	return static_cast<BYTE>(MapVirtualKeyW(static_cast<UINT>(vk), MAPVK_VK_TO_VSC) & 0xFF);
}

using DI_GetDeviceStateFn = HRESULT(STDMETHODCALLTYPE*)(IDirectInputDevice8W*, DWORD, LPVOID);

HRESULT STDMETHODCALLTYPE hk_DI_GetDeviceState(IDirectInputDevice8W* self, DWORD cb, LPVOID data)
{
	DetourGate::Guard guard(g_gate);
	const HRESULT hr = g_vh_di_getstate.original<DI_GetDeviceStateFn>()(self, cb, data);
	if (FAILED(hr) || data == nullptr || !g_active.load(std::memory_order_relaxed))
	{
		return hr;
	}
	hook_note_call(g_id_di_getstate);
	if (cb == 256) // keyboard: BYTE[256] indexed by DIK (scan code); high bit = pressed
	{
		BYTE* keys = static_cast<BYTE*>(data);
		for (int vk = 0; vk < 256; ++vk)
		{
			if (g_key_down[vk].load(std::memory_order_relaxed))
			{
				const BYTE dik = vk_to_dik(vk);
				if (dik != 0)
				{
					keys[dik] |= 0x80;
				}
			}
		}
	}
	else if (cb == sizeof(DIMOUSESTATE) || cb == sizeof(DIMOUSESTATE2)) // mouse (DIMOUSESTATE2 is a superset)
	{
		auto* m = static_cast<DIMOUSESTATE*>(data); // the shared lead fields (lX/lY/lZ/rgbButtons)
		if (g_have_cursor.load(std::memory_order_relaxed))
		{
			const long x = g_cursor_x.load(std::memory_order_relaxed);
			const long y = g_cursor_y.load(std::memory_order_relaxed);
			if (g_di_mouse_primed.load(std::memory_order_relaxed)) // relative delta from our cursor
			{
				m->lX += x - g_di_mouse_last_x.load(std::memory_order_relaxed);
				m->lY += y - g_di_mouse_last_y.load(std::memory_order_relaxed);
			}
			g_di_mouse_last_x.store(x, std::memory_order_relaxed);
			g_di_mouse_last_y.store(y, std::memory_order_relaxed);
			g_di_mouse_primed.store(true, std::memory_order_relaxed);
		}
		if (g_key_down[VK_LBUTTON].load(std::memory_order_relaxed))
		{
			m->rgbButtons[0] |= 0x80;
		}
		if (g_key_down[VK_RBUTTON].load(std::memory_order_relaxed))
		{
			m->rgbButtons[1] |= 0x80;
		}
		if (g_key_down[VK_MBUTTON].load(std::memory_order_relaxed))
		{
			m->rgbButtons[2] |= 0x80;
		}
	}
	return hr;
}

// --- Raw Input GetRawInputData detour ---------------------------------------

// True if `h` is one of our synthetic RAWINPUT handles (an address inside g_raw_slots).
bool is_our_raw(HRAWINPUT h)
{
	auto* p = reinterpret_cast<RAWINPUT*>(h);
	return p >= &g_raw_slots[0] && p < &g_raw_slots[kRawSlots];
}

UINT WINAPI hk_GetRawInputData(HRAWINPUT hri, UINT cmd, LPVOID pData, PUINT pcbSize, UINT cbHeader)
{
	DetourGate::Guard guard(g_gate);
	if (g_active.load(std::memory_order_relaxed) && is_our_raw(hri))
	{
		hook_note_call(g_id_rawinput);
		const RAWINPUT* ri = reinterpret_cast<const RAWINPUT*>(hri);
		const UINT body = ri->header.dwType == RIM_TYPEMOUSE ? sizeof(RAWMOUSE) : sizeof(RAWKEYBOARD);
		const UINT full = sizeof(RAWINPUTHEADER) + body;
		if (pcbSize == nullptr)
		{
			return static_cast<UINT>(-1);
		}
		if (cmd == RID_HEADER)
		{
			if (pData == nullptr)
			{
				*pcbSize = sizeof(RAWINPUTHEADER);
				return 0;
			}
			if (*pcbSize < sizeof(RAWINPUTHEADER))
			{
				return static_cast<UINT>(-1);
			}
			memcpy(pData, &ri->header, sizeof(RAWINPUTHEADER));
			return sizeof(RAWINPUTHEADER);
		}
		// RID_INPUT: the full header + body.
		if (pData == nullptr)
		{
			*pcbSize = full;
			return 0;
		}
		if (*pcbSize < full)
		{
			return static_cast<UINT>(-1);
		}
		memcpy(pData, ri, full);
		return full;
	}
	return g_hk_getrawinputdata.stdcall<UINT>(hri, cmd, pData, pcbSize, cbHeader);
}

// Post a synthetic Raw Input event to `hwnd` (a WM_INPUT carrying one of our g_raw_slots).
void post_raw_key(HWND hwnd, UINT vk, bool down)
{
	if (hwnd == nullptr || !g_hk_getrawinputdata)
	{
		return;
	}
	RAWINPUT& ri = g_raw_slots[g_raw_head.fetch_add(1, std::memory_order_relaxed) % kRawSlots];
	ZeroMemory(&ri, sizeof(ri));
	ri.header.dwType = RIM_TYPEKEYBOARD;
	ri.header.dwSize = sizeof(RAWINPUTHEADER) + sizeof(RAWKEYBOARD);
	ri.data.keyboard.MakeCode = static_cast<USHORT>(MapVirtualKeyW(vk, MAPVK_VK_TO_VSC));
	ri.data.keyboard.Flags = static_cast<USHORT>(down ? RI_KEY_MAKE : RI_KEY_BREAK);
	ri.data.keyboard.VKey = static_cast<USHORT>(vk);
	ri.data.keyboard.Message = down ? WM_KEYDOWN : WM_KEYUP;
	PostMessageW(hwnd, WM_INPUT, RIM_INPUT, reinterpret_cast<LPARAM>(&ri));
}

void post_raw_mouse(HWND hwnd, USHORT button_flags)
{
	if (hwnd == nullptr || !g_hk_getrawinputdata)
	{
		return;
	}
	RAWINPUT& ri = g_raw_slots[g_raw_head.fetch_add(1, std::memory_order_relaxed) % kRawSlots];
	ZeroMemory(&ri, sizeof(ri));
	ri.header.dwType = RIM_TYPEMOUSE;
	ri.header.dwSize = sizeof(RAWINPUTHEADER) + sizeof(RAWMOUSE);
	ri.data.mouse.usFlags = MOUSE_MOVE_RELATIVE;
	ri.data.mouse.usButtonFlags = button_flags;
	PostMessageW(hwnd, WM_INPUT, RIM_INPUT, reinterpret_cast<LPARAM>(&ri));
}

// --- 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));
	// Also feed Raw Input games (button event; relative move isn't in the MKB event stream).
	USHORT rflags = 0;
	if (ev.code == 0)
	{
		rflags = down ? RI_MOUSE_LEFT_BUTTON_DOWN : RI_MOUSE_LEFT_BUTTON_UP;
	}
	else if (ev.code == 1)
	{
		rflags = down ? RI_MOUSE_RIGHT_BUTTON_DOWN : RI_MOUSE_RIGHT_BUTTON_UP;
	}
	else
	{
		rflags = down ? RI_MOUSE_MIDDLE_BUTTON_DOWN : RI_MOUSE_MIDDLE_BUTTON_UP;
	}
	post_raw_mouse(hwnd, rflags);
}

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

// Install the DirectInput GetDeviceState hook: build a kept-alive probe device (its vtable is
// shared by every DI device, including ones the game made before we injected) and swap the slot.
// Returns false (and is retried from mkb_pump) until dinput8.dll is loaded. COM must be initialized
// on the calling thread (the worker thread is).
bool install_dinput_hook()
{
	if (g_vh_di_getstate)
	{
		return true;
	}
	HMODULE di = GetModuleHandleW(L"dinput8.dll");
	if (di == nullptr)
	{
		return false; // not a DirectInput game (yet)
	}
	using PFN_DI8Create = HRESULT(WINAPI*)(HINSTANCE, DWORD, REFIID, LPVOID*, LPUNKNOWN);
	auto create = reinterpret_cast<PFN_DI8Create>(GetProcAddress(di, "DirectInput8Create"));
	if (create == nullptr)
	{
		return false;
	}
	if (g_di_probe == nullptr)
	{
		if (FAILED(create(GetModuleHandleW(nullptr), DIRECTINPUT_VERSION, IID_IDirectInput8W,
		                  reinterpret_cast<void**>(&g_di_probe), nullptr)) ||
		    g_di_probe == nullptr)
		{
			return false;
		}
	}
	if (g_di_probe_kbd == nullptr)
	{
		if (FAILED(g_di_probe->CreateDevice(GUID_SysKeyboard, &g_di_probe_kbd, nullptr)) ||
		    g_di_probe_kbd == nullptr)
		{
			return false;
		}
	}
	const bool ok = g_vh_di_getstate.install(g_di_probe_kbd, kIdx_IDirectInputDevice8_GetDeviceState,
	                                         reinterpret_cast<void*>(&hk_DI_GetDeviceState));
	hook_set_installed(g_id_di_getstate, ok);
	return ok;
}

} // 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);
		g_id_di_getstate = hook_register("IDirectInputDevice8::GetDeviceState", HookSubsys_Mkb);
		g_id_rawinput = hook_register("GetRawInputData", 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);
	// Raw Input: synthesize WM_INPUT (in mkb_pump) + serve it from this hook, for games that read
	// keyboard/mouse via GetRawInputData. GetRawInputData has a clean prologue -> inline hook is OK.
	install_user32_hook(user32, "GetRawInputData", reinterpret_cast<void*>(&hk_GetRawInputData),
	                    g_hk_getrawinputdata, g_id_rawinput);
	// DirectInput: vtable-swap GetDeviceState (best-effort -- dinput8.dll may load later, retried
	// from mkb_pump). The probe is built once and kept alive (avoids COM churn on a re-enable).
	g_di_mouse_primed.store(false, std::memory_order_relaxed);
	install_dinput_hook();

	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 -> no new detour starts
	g_hk_kbstate = {};
	g_hk_cursor = {};
	g_hk_getrawinputdata = {};   // restore GetRawInputData
	g_vh_di_getstate.remove();   // restore the DI GetDeviceState slot (probe kept alive for re-enable)
	g_gate.drain();    // wait for any in-flight polling / DI / raw detour before clearing state
	hook_set_installed(g_id_di_getstate, false);
	hook_set_installed(g_id_rawinput, false);
	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;
	}
	if (!g_vh_di_getstate)
	{
		install_dinput_hook(); // dinput8.dll can load after we installed; keep retrying cheaply
	}

	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));
			}
			post_raw_key(hwnd, ev.code, true); // also feed Raw Input games
			break;
		case Mkb_KeyUp:
			set_key(ev.code, false);
			if (hwnd != nullptr)
			{
				PostMessageW(hwnd, WM_KEYUP, ev.code, key_lparam(ev.code, true));
			}
			post_raw_key(hwnd, ev.code, false);
			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