#include "inject/mkb_forward.hpp"

#include <imgui.h>

#include "coop/protocol.hpp"
#include "injection_panel.hpp"
#include "inject/mkb_map.hpp"

namespace coop {

namespace {

// Map an ImGui key to a Win32 virtual-key. Returns 0 for keys we don't forward.
int imgui_key_to_vk(ImGuiKey k)
{
	if (k >= ImGuiKey_A && k <= ImGuiKey_Z) {
		return 'A' + (k - ImGuiKey_A);
	}
	if (k >= ImGuiKey_0 && k <= ImGuiKey_9) {
		return '0' + (k - ImGuiKey_0);
	}
	if (k >= ImGuiKey_Keypad0 && k <= ImGuiKey_Keypad9) {
		return VK_NUMPAD0 + (k - ImGuiKey_Keypad0);
	}
	if (k >= ImGuiKey_F1 && k <= ImGuiKey_F12) {
		return VK_F1 + (k - ImGuiKey_F1);
	}
	switch (k) {
	case ImGuiKey_Tab:
		return VK_TAB;
	case ImGuiKey_LeftArrow:
		return VK_LEFT;
	case ImGuiKey_RightArrow:
		return VK_RIGHT;
	case ImGuiKey_UpArrow:
		return VK_UP;
	case ImGuiKey_DownArrow:
		return VK_DOWN;
	case ImGuiKey_PageUp:
		return VK_PRIOR;
	case ImGuiKey_PageDown:
		return VK_NEXT;
	case ImGuiKey_Home:
		return VK_HOME;
	case ImGuiKey_End:
		return VK_END;
	case ImGuiKey_Insert:
		return VK_INSERT;
	case ImGuiKey_Delete:
		return VK_DELETE;
	case ImGuiKey_Backspace:
		return VK_BACK;
	case ImGuiKey_Space:
		return VK_SPACE;
	case ImGuiKey_Enter:
		return VK_RETURN;
	case ImGuiKey_Escape:
		return VK_ESCAPE;
	case ImGuiKey_LeftCtrl:
		return VK_LCONTROL;
	case ImGuiKey_LeftShift:
		return VK_LSHIFT;
	case ImGuiKey_LeftAlt:
		return VK_LMENU;
	case ImGuiKey_LeftSuper:
		return VK_LWIN;
	case ImGuiKey_RightCtrl:
		return VK_RCONTROL;
	case ImGuiKey_RightShift:
		return VK_RSHIFT;
	case ImGuiKey_RightAlt:
		return VK_RMENU;
	case ImGuiKey_RightSuper:
		return VK_RWIN;
	case ImGuiKey_Menu:
		return VK_APPS;
	case ImGuiKey_Apostrophe:
		return VK_OEM_7;
	case ImGuiKey_Comma:
		return VK_OEM_COMMA;
	case ImGuiKey_Minus:
		return VK_OEM_MINUS;
	case ImGuiKey_Period:
		return VK_OEM_PERIOD;
	case ImGuiKey_Slash:
		return VK_OEM_2;
	case ImGuiKey_Semicolon:
		return VK_OEM_1;
	case ImGuiKey_Equal:
		return VK_OEM_PLUS;
	case ImGuiKey_LeftBracket:
		return VK_OEM_4;
	case ImGuiKey_Backslash:
		return VK_OEM_5;
	case ImGuiKey_RightBracket:
		return VK_OEM_6;
	case ImGuiKey_GraveAccent:
		return VK_OEM_3;
	case ImGuiKey_CapsLock:
		return VK_CAPITAL;
	case ImGuiKey_ScrollLock:
		return VK_SCROLL;
	case ImGuiKey_NumLock:
		return VK_NUMLOCK;
	case ImGuiKey_PrintScreen:
		return VK_SNAPSHOT;
	case ImGuiKey_Pause:
		return VK_PAUSE;
	case ImGuiKey_KeypadDecimal:
		return VK_DECIMAL;
	case ImGuiKey_KeypadDivide:
		return VK_DIVIDE;
	case ImGuiKey_KeypadMultiply:
		return VK_MULTIPLY;
	case ImGuiKey_KeypadSubtract:
		return VK_SUBTRACT;
	case ImGuiKey_KeypadAdd:
		return VK_ADD;
	case ImGuiKey_KeypadEnter:
		return VK_RETURN;
	default:
		return 0;
	}
}

// Last valid game-client position, so a button release that lands on a letterbox
// bar still goes to the game (no stuck buttons).
int g_last_gx = 0;
int g_last_gy = 0;

// What we've forwarded as currently-held, so we can always release it -- even when forwarding stops
// (mirror off, focus lost, ImGui grabbed the input, the game died) -- so nothing sticks in the guest.
bool g_mouse_down[3] = {};
bool g_key_down[256] = {}; // indexed by VK

void release_held_keys(InjectionPanel& injection)
{
	for (int vk = 0; vk < 256; ++vk) {
		if (g_key_down[vk]) {
			injection.push_mkb(MkbEvent{Mkb_KeyUp, static_cast<std::uint32_t>(vk), 0, 0});
			g_key_down[vk] = false;
		}
	}
}

void release_held_mouse(InjectionPanel& injection)
{
	for (int b = 0; b < 3; ++b) {
		if (g_mouse_down[b]) {
			injection.push_mkb(MkbEvent{Mkb_MouseUp, static_cast<std::uint32_t>(b), g_last_gx, g_last_gy});
			g_mouse_down[b] = false;
		}
	}
}

} // namespace

void forward_mkb_frame(InjectionPanel& injection, HWND host_hwnd, bool mirroring, bool source_hooked)
{
	const HWND game = injection.game_hwnd();
	// Under RPT the guest's MKB lands on our (focused) window; only forward then, so the operator's
	// own desktop use isn't injected. If we can't forward for ANY reason -- subsystem off, we lost
	// focus, or the game is gone -- release everything we're still holding first, so a key/button
	// held at that moment doesn't stick down in the guest.
	if (!injection.mkb_enabled() || GetForegroundWindow() != host_hwnd || game == nullptr || !IsWindow(game)) {
		release_held_keys(injection);
		release_held_mouse(injection);
		return;
	}

	ImGuiIO& io = ImGui::GetIO();

	// --- Keyboard (unless ImGui is using it for e.g. a text field -- then release what we hold) ---
	if (io.WantCaptureKeyboard) {
		release_held_keys(injection);
	} else {
		for (ImGuiKey k = ImGuiKey_NamedKey_BEGIN; k < ImGuiKey_NamedKey_END; k = static_cast<ImGuiKey>(k + 1)) {
			const int vk = imgui_key_to_vk(k);
			if (vk == 0) {
				continue;
			}
			if (ImGui::IsKeyPressed(k, false)) {
				injection.push_mkb(MkbEvent{Mkb_KeyDown, static_cast<std::uint32_t>(vk), 0, 0});
				g_key_down[vk & 0xFF] = true;
			}
			if (ImGui::IsKeyReleased(k)) {
				injection.push_mkb(MkbEvent{Mkb_KeyUp, static_cast<std::uint32_t>(vk), 0, 0});
				g_key_down[vk & 0xFF] = false;
			}
		}
		for (int i = 0; i < io.InputQueueCharacters.Size; ++i) {
			const ImWchar c = io.InputQueueCharacters[i];
			if (c != 0) {
				injection.push_mkb(MkbEvent{Mkb_Char, static_cast<std::uint32_t>(c), 0, 0});
			}
		}
	}

	// --- Mouse (clicks + wheel only, and only while mirroring and ImGui isn't using the mouse) ---
	// When we're not forwarding the mouse, still release any button we hold (below), so it can't stick.
	const bool forwarding_mouse = mirroring && !io.WantCaptureMouse;
	if (!forwarding_mouse) {
		release_held_mouse(injection);
		return;
	}

	RECT host_client{};
	GetClientRect(host_hwnd, &host_client);

	MkbMapInput m;
	m.host_x = static_cast<int>(io.MousePos.x);
	m.host_y = static_cast<int>(io.MousePos.y);
	m.dst_w = host_client.right;
	m.dst_h = host_client.bottom;
	if (source_hooked) {
		// Hooked capture mirrors the backbuffer (client area), no decorations.
		const VideoShareView v = injection.video_share();
		m.src_w = m.client_w = static_cast<int>(v.width);
		m.src_h = m.client_h = static_cast<int>(v.height);
	} else {
		// WGC captures the whole window; the client area sits at a decoration offset.
		RECT wr{}, cr{};
		POINT client_origin{0, 0};
		GetWindowRect(game, &wr);
		GetClientRect(game, &cr);
		ClientToScreen(game, &client_origin);
		m.src_w = wr.right - wr.left;
		m.src_h = wr.bottom - wr.top;
		m.client_w = cr.right;
		m.client_h = cr.bottom;
		m.client_off_x = client_origin.x - wr.left;
		m.client_off_y = client_origin.y - wr.top;
	}

	int gx = 0, gy = 0;
	const bool on_game = map_host_to_game_client(m, gx, gy);
	if (on_game) {
		g_last_gx = gx;
		g_last_gy = gy;
	}
	const int mx = on_game ? gx : g_last_gx;
	const int my = on_game ? gy : g_last_gy;

	for (int button = 0; button < 3; ++button) // 0=left, 1=right, 2=middle
	{
		if (on_game && ImGui::IsMouseClicked(button)) {
			injection.push_mkb(MkbEvent{Mkb_MouseDown, static_cast<std::uint32_t>(button), mx, my});
			g_mouse_down[button] = true;
		}
		if (ImGui::IsMouseReleased(button)) // send the up even off-game, so nothing sticks
		{
			injection.push_mkb(MkbEvent{Mkb_MouseUp, static_cast<std::uint32_t>(button), mx, my});
			g_mouse_down[button] = false;
		}
	}
	if (on_game && io.MouseWheel != 0.0f) {
		const int delta = static_cast<int>(io.MouseWheel * WHEEL_DELTA);
		injection.push_mkb(MkbEvent{Mkb_Wheel, static_cast<std::uint32_t>(delta), mx, my});
	}
}

} // namespace coop