#include "d3d9_hook.hpp"

#include <atomic>
#include <vector>

#include <windows.h>

#include <d3d9.h>
#include <d3d11.h>
#include <dxgi1_2.h>

#include <safetyhook.hpp>

#include "coop/protocol.hpp"
#include "coop/shared_memory.hpp"
#include "debug_log.hpp"
#include "hook_guard.hpp"
#include "hook_install.hpp"
#include "hook_registry.hpp"

namespace coop::hook
{

namespace
{

DetourGate g_gate; // drains in-flight Present detours before remove frees the shared D3D state

// IDirect3DDevice9 vtable: IUnknown 0-2, then the device methods. Present is index 17
// (TestCooperativeLevel 3, GetAvailableTextureMem 4, EvictManagedResources 5, GetDirect3D 6,
// GetDeviceCaps 7, GetDisplayMode 8, GetCreationParameters 9, SetCursorProperties 10,
// SetCursorPosition 11, ShowCursor 12, CreateAdditionalSwapChain 13, GetSwapChain 14,
// GetNumberOfSwapChains 15, Reset 16, Present 17). IDirect3DDevice9Ex games that call the
// inherited Present hit this too (PresentEx is a separate, higher slot -- not needed for the
// common case / the mock).
constexpr unsigned kIdx_IDirect3DDevice9_Present = 17;

IpcClient* g_ipc = nullptr;
unsigned long g_pid = 0;

safetyhook::InlineHook g_hk_present9;
int g_id_present9 = -1;

std::atomic<std::uint64_t> g_presents{0};
std::atomic<std::uint64_t> g_frames_shared{0};
bool g_unsupported_logged = false;

// Our own D3D11 device hosting the shared texture (the D3D9 game has no D3D11 device).
ID3D11Device* g_device = nullptr;
ID3D11DeviceContext* g_ctx = nullptr;
ID3D11Texture2D* g_shared_tex = nullptr;
IDXGIKeyedMutex* g_shared_mutex = nullptr;
HANDLE g_shared_handle = nullptr;
UINT g_share_w = 0;
UINT g_share_h = 0;

// System-memory read-back surface on the game's D3D9 device (GetRenderTargetData target).
IDirect3DSurface9* g_sysmem = nullptr;
IDirect3DDevice9* g_sysmem_dev = nullptr;
UINT g_sysmem_w = 0;
UINT g_sysmem_h = 0;
D3DFORMAT g_sysmem_fmt = D3DFMT_UNKNOWN;

std::vector<unsigned char> g_rgba; // swizzled RGBA, uploaded to D3D11

// Present may be issued re-entrantly by some engines; capture only on the outermost call.
thread_local bool t_in_present = false;

void* vtable_method(void* obj, unsigned index)
{
	return (*reinterpret_cast<void***>(obj))[index];
}

bool ensure_device()
{
	if (g_device != nullptr)
	{
		return true;
	}
	const HRESULT hr = D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, 0, nullptr, 0,
	                                     D3D11_SDK_VERSION, &g_device, nullptr, &g_ctx);
	if (FAILED(hr) || g_device == nullptr)
	{
		logf("d3d9: D3D11CreateDevice failed hr=0x%08lX", static_cast<unsigned long>(hr));
		return false;
	}
	return true;
}

void release_shared()
{
	if (g_shared_mutex != nullptr)
	{
		g_shared_mutex->Release();
		g_shared_mutex = nullptr;
	}
	if (g_shared_tex != nullptr)
	{
		g_shared_tex->Release();
		g_shared_tex = nullptr;
	}
	if (g_shared_handle != nullptr)
	{
		CloseHandle(g_shared_handle);
		g_shared_handle = nullptr;
	}
	g_share_w = g_share_h = 0;
}

void release_sysmem()
{
	if (g_sysmem != nullptr)
	{
		g_sysmem->Release();
		g_sysmem = nullptr;
	}
	if (g_sysmem_dev != nullptr)
	{
		g_sysmem_dev->Release();
		g_sysmem_dev = nullptr;
	}
	g_sysmem_w = g_sysmem_h = 0;
	g_sysmem_fmt = D3DFMT_UNKNOWN;
}

bool ensure_shared_texture(UINT w, UINT h)
{
	if (g_shared_tex != nullptr && g_share_w == w && g_share_h == h)
	{
		return true;
	}
	release_shared();

	D3D11_TEXTURE2D_DESC desc{};
	desc.Width = w;
	desc.Height = h;
	desc.MipLevels = 1;
	desc.ArraySize = 1;
	desc.Format = DXGI_FORMAT_R8G8B8A8_UNORM; // we swizzle the D3D9 BGRA backbuffer to RGBA
	desc.SampleDesc.Count = 1;
	desc.Usage = D3D11_USAGE_DEFAULT;
	desc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
	desc.MiscFlags = D3D11_RESOURCE_MISC_SHARED_NTHANDLE | D3D11_RESOURCE_MISC_SHARED_KEYEDMUTEX;

	if (FAILED(g_device->CreateTexture2D(&desc, nullptr, &g_shared_tex)) || g_shared_tex == nullptr)
	{
		return false;
	}
	IDXGIResource1* res = nullptr;
	if (FAILED(g_shared_tex->QueryInterface(__uuidof(IDXGIResource1), reinterpret_cast<void**>(&res))) ||
	    res == nullptr)
	{
		release_shared();
		return false;
	}
	const std::wstring name = video_share_name(g_pid);
	const HRESULT hr = res->CreateSharedHandle(
	    nullptr, DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE, name.c_str(), &g_shared_handle);
	res->Release();
	if (FAILED(hr) || g_shared_handle == nullptr)
	{
		release_shared();
		return false;
	}
	if (FAILED(g_shared_tex->QueryInterface(__uuidof(IDXGIKeyedMutex), reinterpret_cast<void**>(&g_shared_mutex))))
	{
		release_shared();
		return false;
	}
	g_share_w = w;
	g_share_h = h;
	logf("d3d9: shared texture ready %ux%u name=%ls", w, h, name.c_str());
	return true;
}

// Read the game's D3D9 backbuffer back to system memory, swizzle BGRA->RGBA, and upload it.
void capture_d3d9(IDirect3DDevice9* dev)
{
	IDirect3DSurface9* back = nullptr;
	if (FAILED(dev->GetBackBuffer(0, 0, D3DBACKBUFFER_TYPE_MONO, &back)) || back == nullptr)
	{
		return;
	}
	D3DSURFACE_DESC d{};
	back->GetDesc(&d);
	const UINT w = d.Width;
	const UINT h = d.Height;
	// We only handle the standard 32-bit BGRX/BGRA back buffers (the common D3D9 case).
	if ((d.Format != D3DFMT_X8R8G8B8 && d.Format != D3DFMT_A8R8G8B8) || w == 0 || h == 0)
	{
		if (!g_unsupported_logged)
		{
			logf("d3d9: unsupported backbuffer format=%d (only X8R8G8B8 / A8R8G8B8); idle", static_cast<int>(d.Format));
			g_unsupported_logged = true;
		}
		back->Release();
		return;
	}

	// (Re)create the system-memory read-back surface on the game's device.
	if (!(g_sysmem != nullptr && g_sysmem_dev == dev && g_sysmem_w == w && g_sysmem_h == h && g_sysmem_fmt == d.Format))
	{
		release_sysmem();
		if (SUCCEEDED(dev->CreateOffscreenPlainSurface(w, h, d.Format, D3DPOOL_SYSTEMMEM, &g_sysmem, nullptr)) &&
		    g_sysmem != nullptr)
		{
			g_sysmem_dev = dev;
			dev->AddRef();
			g_sysmem_w = w;
			g_sysmem_h = h;
			g_sysmem_fmt = d.Format;
		}
	}

	bool shared = false;
	if (g_sysmem != nullptr && SUCCEEDED(dev->GetRenderTargetData(back, g_sysmem))) // GPU->sysmem, blocks
	{
		D3DLOCKED_RECT lr{};
		if (SUCCEEDED(g_sysmem->LockRect(&lr, nullptr, D3DLOCK_READONLY)) && lr.pBits != nullptr)
		{
			const size_t dst_row = static_cast<size_t>(w) * 4;
			if (g_rgba.size() != dst_row * h)
			{
				g_rgba.resize(dst_row * h);
			}
			// X8R8G8B8 / A8R8G8B8 store as little-endian 0xAARRGGBB -> bytes B,G,R,A. Swizzle to
			// R,G,B,A and force opaque alpha so the host's RGBA decode matches the other backends.
			for (UINT y = 0; y < h; ++y)
			{
				const unsigned char* src = static_cast<const unsigned char*>(lr.pBits) + static_cast<size_t>(y) * lr.Pitch;
				unsigned char* out = g_rgba.data() + static_cast<size_t>(y) * dst_row;
				for (UINT x = 0; x < w; ++x)
				{
					out[x * 4 + 0] = src[x * 4 + 2]; // R
					out[x * 4 + 1] = src[x * 4 + 1]; // G
					out[x * 4 + 2] = src[x * 4 + 0]; // B
					out[x * 4 + 3] = 255;            // A
				}
			}
			g_sysmem->UnlockRect();

			if (ensure_device() && ensure_shared_texture(w, h) && g_shared_mutex != nullptr &&
			    g_shared_mutex->AcquireSync(kVideoMutexKey, 8) == S_OK)
			{
				g_ctx->UpdateSubresource(g_shared_tex, 0, nullptr, g_rgba.data(), static_cast<UINT>(dst_row), 0);
				g_ctx->Flush();
				g_shared_mutex->ReleaseSync(kVideoMutexKey);
				shared = true;
			}
		}
	}

	if (shared)
	{
		g_frames_shared.fetch_add(1, std::memory_order_relaxed);
		if (g_ipc != nullptr)
		{
			g_ipc->publish_video_frame(w, h, static_cast<std::uint32_t>(DXGI_FORMAT_R8G8B8A8_UNORM));
		}
	}
	back->Release();
}

HRESULT STDMETHODCALLTYPE hk_Present9(IDirect3DDevice9* dev, const RECT* src, const RECT* dst, HWND wnd,
                                      const RGNDATA* dirty)
{
	DetourGate::Guard guard(g_gate); // keep the shared D3D state alive for this whole detour
	hook_note_call(g_id_present9);
	g_presents.fetch_add(1, std::memory_order_relaxed);
	if (g_ipc != nullptr)
	{
		g_ipc->note_present();
	}
	if (!t_in_present)
	{
		t_in_present = true;
		capture_d3d9(dev);
		t_in_present = false;
	}
	// stdcall(): IDirect3DDevice9::Present is __stdcall; call() would invoke the trampoline as
	// __cdecl on x86 -> ESP imbalance -> crash. No-op on x64. (Present has a clean prologue, so
	// an inline hook is safe -- unlike the WASAPI COM methods, which need a vtable swap.)
	return g_hk_present9.stdcall<HRESULT>(dev, src, dst, wnd, dirty);
}

// Create a throwaway D3D9 device to read IDirect3DDevice9::Present's address, so we can
// inline-hook it (catching the game's existing device regardless of when we injected). Returns
// null when d3d9.dll isn't loaded (not a D3D9 game) or a probe device can't be created.
void* grab_present9_address()
{
	HMODULE d3d9 = GetModuleHandleW(L"d3d9.dll");
	if (d3d9 == nullptr)
	{
		return nullptr; // not a D3D9 game
	}
	using PFN_Direct3DCreate9 = IDirect3D9*(WINAPI*)(UINT);
	auto create = reinterpret_cast<PFN_Direct3DCreate9>(GetProcAddress(d3d9, "Direct3DCreate9"));
	if (create == nullptr)
	{
		return nullptr;
	}
	IDirect3D9* d3d = create(D3D_SDK_VERSION);
	if (d3d == nullptr)
	{
		return nullptr;
	}

	WNDCLASSEXW wc{};
	wc.cbSize = sizeof(wc);
	wc.lpfnWndProc = DefWindowProcW;
	wc.hInstance = GetModuleHandleW(nullptr);
	wc.lpszClassName = L"coop_d3d9_probe";
	RegisterClassExW(&wc);
	HWND hwnd = CreateWindowExW(0, wc.lpszClassName, L"", WS_OVERLAPPEDWINDOW, 0, 0, 8, 8, nullptr, nullptr,
	                            wc.hInstance, nullptr);

	void* present = nullptr;
	if (hwnd != nullptr)
	{
		D3DPRESENT_PARAMETERS pp{};
		pp.BackBufferWidth = 8;
		pp.BackBufferHeight = 8;
		pp.BackBufferFormat = D3DFMT_X8R8G8B8;
		pp.BackBufferCount = 1;
		pp.SwapEffect = D3DSWAPEFFECT_DISCARD;
		pp.hDeviceWindow = hwnd;
		pp.Windowed = TRUE;
		IDirect3DDevice9* dev = nullptr;
		if (SUCCEEDED(d3d->CreateDevice(D3DADAPTER_DEFAULT, D3DDEVTYPE_HAL, hwnd,
		                                D3DCREATE_HARDWARE_VERTEXPROCESSING | D3DCREATE_MULTITHREADED, &pp, &dev)) &&
		    dev != nullptr)
		{
			present = vtable_method(dev, kIdx_IDirect3DDevice9_Present);
			dev->Release();
		}
		DestroyWindow(hwnd);
	}
	UnregisterClassW(wc.lpszClassName, wc.hInstance);
	d3d->Release();
	return present;
}

} // namespace

bool install_d3d9_hooks(IpcClient& ipc)
{
	g_ipc = &ipc;
	g_pid = GetCurrentProcessId();
	if (g_hk_present9.enabled())
	{
		return true; // already installed (persistent hook; re-install below re-enables it)
	}
	g_id_present9 = hook_register("IDirect3DDevice9::Present", HookSubsys_Video);
	g_unsupported_logged = false;

	void* present = grab_present9_address();
	if (present == nullptr)
	{
		hook_set_installed(g_id_present9, false); // not a D3D9 game (or no probe device)
		return false;
	}
	install_inline(g_hk_present9, present, &hk_Present9);
	hook_set_installed(g_id_present9, static_cast<bool>(g_hk_present9));
	logf("install_d3d9_hooks: Present=%p hooked=%d", present, static_cast<bool>(g_hk_present9) ? 1 : 0);
	return static_cast<bool>(g_hk_present9);
}

void remove_d3d9_hooks()
{
	// DISABLE (persistent model -- never destroy during the session): restores Present's bytes under
	// thread suspension (no new detour) but keeps the trampoline alive, so an in-flight detour about
	// to call g_hk_present9.stdcall() (the trampoline) never has it freed under it. Destroying (= {})
	// would free it -- a UAF the thousands/s storm hits reliably (0xC0000005). Disable -> drain ->
	// leave alive (re-install re-enables; see hook_install.hpp).
	disable_for_removal(g_hk_present9);
	hook_set_installed(g_id_present9, false);
	g_gate.drain();
	// Persistent hook: keep g_hk_present9 ALIVE (disabled) so a stale detour's trampoline call is
	// never freed -- re-install re-enables it (see hook_install.hpp).
	release_shared();
	release_sysmem();
	if (g_ctx != nullptr)
	{
		g_ctx->Release();
		g_ctx = nullptr;
	}
	if (g_device != nullptr)
	{
		g_device->Release();
		g_device = nullptr;
	}
	g_rgba.clear();
	g_presents.store(0, std::memory_order_relaxed);
	g_frames_shared.store(0, std::memory_order_relaxed);
	g_ipc = nullptr;
}

std::uint64_t d3d9_presents()
{
	return g_presents.load(std::memory_order_relaxed);
}

std::uint64_t d3d9_frames_shared()
{
	return g_frames_shared.load(std::memory_order_relaxed);
}

} // namespace coop::hook