#include "xinput_hook.hpp"

#include <array>
#include <atomic>
#include <vector>

#include <windows.h>
#include <xinput.h>

#include <safetyhook.hpp>

namespace coop::hook
{

namespace
{

// XInput guide-button bit, reported only by the undocumented ordinal-100
// XInputGetStateEx that many games use. Mirrors how Steam/x360ce expose it.
constexpr std::uint16_t kGuideButton = 0x0400;

const IpcClient* g_ipc = nullptr;
std::vector<safetyhook::InlineHook> g_hooks;

// Last good snapshot, so a momentary failed IPC read (host mid-write) doesn't
// flicker the controller as disconnected inside the game.
std::array<CoopPadState, kMaxPads> g_cache;

void refresh_cache()
{
	if (g_ipc == nullptr)
	{
		return;
	}
	CoopPadState pads[kMaxPads];
	std::uint32_t count = 0;
	if (g_ipc->snapshot(pads, count))
	{
		for (std::uint32_t i = 0; i < kMaxPads; ++i)
		{
			g_cache[i] = pads[i];
		}
	}
}

void fill_gamepad(const CoopPadState& pad, XINPUT_GAMEPAD& out)
{
	out.wButtons = pad.buttons;
	out.bLeftTrigger = pad.left_trigger;
	out.bRightTrigger = pad.right_trigger;
	out.sThumbLX = pad.thumb_lx;
	out.sThumbLY = pad.thumb_ly;
	out.sThumbRX = pad.thumb_rx;
	out.sThumbRY = pad.thumb_ry;
}

// Core of every state query. `keep_guide` drops the guide bit for the plain
// (documented) XInputGetState, which must not report it.
DWORD query_state(DWORD user_index, XINPUT_STATE* state, bool keep_guide)
{
	if (state == nullptr || user_index >= kMaxPads)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}
	refresh_cache();
	const CoopPadState& pad = g_cache[user_index];
	if (!pad.connected)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}

	XINPUT_STATE result = {};
	result.dwPacketNumber = pad.packet;
	fill_gamepad(pad, result.Gamepad);
	if (!keep_guide)
	{
		result.Gamepad.wButtons &= ~kGuideButton;
	}
	*state = result;
	return ERROR_SUCCESS;
}

DWORD WINAPI hk_XInputGetState(DWORD user_index, XINPUT_STATE* state)
{
	return query_state(user_index, state, /*keep_guide=*/false);
}

DWORD WINAPI hk_XInputGetStateEx(DWORD user_index, XINPUT_STATE* state)
{
	return query_state(user_index, state, /*keep_guide=*/true);
}

DWORD WINAPI hk_XInputGetCapabilities(DWORD user_index, DWORD /*flags*/, XINPUT_CAPABILITIES* caps)
{
	if (caps == nullptr || user_index >= kMaxPads)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}
	refresh_cache();
	if (!g_cache[user_index].connected)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}

	// Advertise a standard wired Xbox-style gamepad with all controls present.
	XINPUT_CAPABILITIES result = {};
	result.Type = XINPUT_DEVTYPE_GAMEPAD;
	result.SubType = XINPUT_DEVSUBTYPE_GAMEPAD;
	result.Flags = 0;
	result.Gamepad.wButtons = 0xF3FF; // all standard buttons reachable
	result.Gamepad.bLeftTrigger = 0xFF;
	result.Gamepad.bRightTrigger = 0xFF;
	result.Gamepad.sThumbLX = static_cast<SHORT>(0x7FFF);
	result.Gamepad.sThumbLY = static_cast<SHORT>(0x7FFF);
	result.Gamepad.sThumbRX = static_cast<SHORT>(0x7FFF);
	result.Gamepad.sThumbRY = static_cast<SHORT>(0x7FFF);
	*caps = result;
	return ERROR_SUCCESS;
}

// Swallow rumble: it would otherwise be sent to whatever physical device sits at
// this index on the host machine. Forwarding it back to the guest is a later
// phase; for now report success so the game's logic is happy.
DWORD WINAPI hk_XInputSetState(DWORD user_index, XINPUT_VIBRATION* /*vibration*/)
{
	if (user_index >= kMaxPads || !g_cache[user_index].connected)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}
	return ERROR_SUCCESS;
}

void hook_export(HMODULE module, const char* name, void* detour)
{
	if (module == nullptr)
	{
		return;
	}
	if (void* target = reinterpret_cast<void*>(GetProcAddress(module, name)))
	{
		g_hooks.emplace_back(safetyhook::create_inline(target, detour));
	}
}

void hook_ordinal(HMODULE module, WORD ordinal, void* detour)
{
	if (module == nullptr)
	{
		return;
	}
	if (void* target = reinterpret_cast<void*>(GetProcAddress(module, MAKEINTRESOURCEA(ordinal))))
	{
		g_hooks.emplace_back(safetyhook::create_inline(target, detour));
	}
}

} // namespace

bool install_xinput_hooks(const IpcClient& ipc)
{
	if (!g_hooks.empty())
	{
		return true; // already installed
	}
	g_ipc = &ipc;
	refresh_cache();

	// A process generally loads exactly one of these, but hook every one that is
	// present so we don't miss the one the game actually calls.
	const wchar_t* modules[] = {L"xinput1_4.dll", L"xinput1_3.dll", L"xinput9_1_0.dll", L"xinputuap.dll"};
	for (const wchar_t* name : modules)
	{
		HMODULE module = GetModuleHandleW(name);
		if (module == nullptr)
		{
			continue;
		}
		hook_export(module, "XInputGetState", reinterpret_cast<void*>(&hk_XInputGetState));
		hook_ordinal(module, 100, reinterpret_cast<void*>(&hk_XInputGetStateEx));
		hook_export(module, "XInputGetCapabilities", reinterpret_cast<void*>(&hk_XInputGetCapabilities));
		hook_export(module, "XInputSetState", reinterpret_cast<void*>(&hk_XInputSetState));
	}
	return !g_hooks.empty();
}

void remove_xinput_hooks()
{
	g_hooks.clear(); // InlineHook destructor restores the original bytes
	g_ipc = nullptr;
}

} // namespace coop::hook