// In-process self-test for the D3D12 path of the Present hook
// (hook/src/present_hook.cpp). This process plays both "game" and "host": it
// installs the Present hook, then drives a real D3D12 swapchain -- clears a
// backbuffer to a known color and calls Present. IDXGISwapChain::Present is the
// same DXGI vtable function for D3D11 and D3D12 swapchains, so the inline hook
// fires; the D3D12 backbuffer can't be a GetBuffer'd ID3D11Texture2D, so the hook
// must bridge it via D3D11On12 and CopyResource it into the shared texture. A
// second D3D11 device then opens that texture by name and verifies the color.
//
// Requires a D3D12-capable GPU; on a machine without one it reports SKIP, exit 0.

#include <cstdint>
#include <cstdio>

#include <windows.h>

#include <d3d11_1.h>
#include <d3d12.h>
#include <dxgi1_4.h>

#include "coop/protocol.hpp"
#include "coop/shared_memory.hpp"
#include "ipc_client.hpp"
#include "present_hook.hpp"
#include "present_overhead.hpp"

using namespace coop;

namespace
{

int g_failures = 0;
void check(bool ok, const char* what)
{
	if (!ok)
	{
		std::printf("  FAIL: %s\n", what);
		++g_failures;
	}
}

template <typename T>
void release(T*& p)
{
	if (p)
	{
		p->Release();
		p = nullptr;
	}
}

bool near_byte(std::uint8_t got, int expected)
{
	const int d = static_cast<int>(got) - expected;
	return d >= -4 && d <= 4;
}

constexpr UINT kW = 256;
constexpr UINT kH = 256;
constexpr UINT kBuffers = 2;
constexpr float kClear[4] = {0.20f, 0.40f, 0.60f, 1.0f}; // ~ {51, 102, 153}

} // namespace

int main()
{
	// --- Host side: shared block named by our pid (the hook opens the same name). ---
	SharedMemory shm;
	if (!shm.create(shared_memory_name(GetCurrentProcessId()), sizeof(SharedBlock)))
	{
		std::printf("FAIL: create shared memory\n");
		return 1;
	}
	auto* block = shm.as<SharedBlock>();
	block->version = kProtocolVersion;
	block->sequence.store(0, std::memory_order_relaxed);
	block->magic = kProtocolMagic;

	// --- D3D12 device + direct queue. SKIP if the machine has no D3D12. ---
	ID3D12Device* device = nullptr;
	if (FAILED(D3D12CreateDevice(nullptr, D3D_FEATURE_LEVEL_11_0, IID_PPV_ARGS(&device))) || device == nullptr)
	{
		std::printf("SKIP: no D3D12 device on this machine\n");
		return 0;
	}
	ID3D12CommandQueue* queue = nullptr;
	D3D12_COMMAND_QUEUE_DESC qd{};
	qd.Type = D3D12_COMMAND_LIST_TYPE_DIRECT;
	check(SUCCEEDED(device->CreateCommandQueue(&qd, IID_PPV_ARGS(&queue))), "create command queue");

	IDXGIFactory4* factory = nullptr;
	check(SUCCEEDED(CreateDXGIFactory1(IID_PPV_ARGS(&factory))), "create DXGI factory");

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

	DXGI_SWAP_CHAIN_DESC1 scd{};
	scd.Width = kW;
	scd.Height = kH;
	scd.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
	scd.SampleDesc.Count = 1;
	scd.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
	scd.BufferCount = kBuffers;
	scd.SwapEffect = DXGI_SWAP_EFFECT_FLIP_DISCARD;

	IDXGISwapChain1* sc1 = nullptr;
	check(SUCCEEDED(factory->CreateSwapChainForHwnd(queue, hwnd, &scd, nullptr, nullptr, &sc1)),
	      "create D3D12 swapchain");
	IDXGISwapChain* swapchain = nullptr;
	IDXGISwapChain3* sc3 = nullptr; // for GetCurrentBackBufferIndex
	if (sc1 != nullptr)
	{
		sc1->QueryInterface(IID_PPV_ARGS(&swapchain));
		sc1->QueryInterface(IID_PPV_ARGS(&sc3));
	}

	// RTV heap + render targets for the swapchain buffers.
	ID3D12DescriptorHeap* rtv_heap = nullptr;
	D3D12_DESCRIPTOR_HEAP_DESC hd{};
	hd.Type = D3D12_DESCRIPTOR_HEAP_TYPE_RTV;
	hd.NumDescriptors = kBuffers;
	device->CreateDescriptorHeap(&hd, IID_PPV_ARGS(&rtv_heap));
	const UINT rtv_size = device->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_RTV);
	ID3D12Resource* render_targets[kBuffers] = {};
	if (swapchain != nullptr && rtv_heap != nullptr)
	{
		D3D12_CPU_DESCRIPTOR_HANDLE rtv = rtv_heap->GetCPUDescriptorHandleForHeapStart();
		for (UINT i = 0; i < kBuffers; ++i)
		{
			swapchain->GetBuffer(i, IID_PPV_ARGS(&render_targets[i]));
			device->CreateRenderTargetView(render_targets[i], nullptr, rtv);
			rtv.ptr += rtv_size;
		}
	}

	ID3D12CommandAllocator* allocator = nullptr;
	device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT, IID_PPV_ARGS(&allocator));
	ID3D12GraphicsCommandList* cmdlist = nullptr;
	device->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, allocator, nullptr, IID_PPV_ARGS(&cmdlist));
	if (cmdlist != nullptr)
	{
		cmdlist->Close();
	}
	ID3D12Fence* fence = nullptr;
	device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence));
	HANDLE fence_event = CreateEventW(nullptr, FALSE, FALSE, nullptr);
	UINT64 fence_value = 0;

	// --- Install the Present hook, then render + present a few frames. ---
	hook::IpcClient ipc;
	check(ipc.connect(10, 5), "IPC client connect");
	check(hook::install_present_hooks(ipc), "install Present hooks");

	const bool can_render =
	    swapchain != nullptr && sc3 != nullptr && allocator != nullptr && cmdlist != nullptr && fence != nullptr;
	for (int frame = 0; frame < 4 && can_render; ++frame)
	{
		const UINT idx = sc3->GetCurrentBackBufferIndex();
		allocator->Reset();
		cmdlist->Reset(allocator, nullptr);

		D3D12_RESOURCE_BARRIER b{};
		b.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
		b.Transition.pResource = render_targets[idx];
		b.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
		b.Transition.StateBefore = D3D12_RESOURCE_STATE_PRESENT;
		b.Transition.StateAfter = D3D12_RESOURCE_STATE_RENDER_TARGET;
		cmdlist->ResourceBarrier(1, &b);

		D3D12_CPU_DESCRIPTOR_HANDLE rtv = rtv_heap->GetCPUDescriptorHandleForHeapStart();
		rtv.ptr += static_cast<SIZE_T>(idx) * rtv_size;
		cmdlist->ClearRenderTargetView(rtv, kClear, 0, nullptr);

		std::swap(b.Transition.StateBefore, b.Transition.StateAfter); // RENDER_TARGET -> PRESENT
		cmdlist->ResourceBarrier(1, &b);
		cmdlist->Close();

		ID3D12CommandList* lists[] = {cmdlist};
		queue->ExecuteCommandLists(1, lists);

		swapchain->Present(0, 0); // -> hooked IDXGISwapChain::Present -> D3D11On12 bridge

		// Block until the GPU finished this frame (keeps the test simple + correct).
		queue->Signal(fence, ++fence_value);
		if (fence->GetCompletedValue() < fence_value)
		{
			fence->SetEventOnCompletion(fence_value, fence_event);
			WaitForSingleObject(fence_event, 1000);
		}
	}

	std::printf("present_calls=%llu frames_shared=%llu video{gen=%u %ux%u fmt=%u}\n",
	            static_cast<unsigned long long>(hook::present_calls()),
	            static_cast<unsigned long long>(hook::present_frames_shared()), block->video.generation.load(),
	            block->video.width, block->video.height, block->video.format);

	if (can_render)
	{
		check(hook::present_calls() >= 3, "Present detour fired for the D3D12 swapchain");
		check(hook::present_frames_shared() > 0, "D3D12 backbuffer bridged into the shared texture");
		check(block->video.generation.load() > 0, "video generation published to IPC");
		check(block->video.width == kW && block->video.height == kH, "shared dimensions published");

		// --- Consumer side: open the shared texture by name and verify the color. ---
		ID3D11Device* devB = nullptr;
		ID3D11DeviceContext* ctxB = nullptr;
		if (SUCCEEDED(D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, 0, nullptr, 0, D3D11_SDK_VERSION,
		                                &devB, nullptr, &ctxB)))
		{
			ID3D11Device1* dev1 = nullptr;
			devB->QueryInterface(IID_PPV_ARGS(&dev1));
			const std::wstring name = video_share_name(GetCurrentProcessId());
			ID3D11Texture2D* sharedB = nullptr;
			IDXGIKeyedMutex* km = nullptr;
			if (dev1 != nullptr &&
			    SUCCEEDED(dev1->OpenSharedResourceByName(name.c_str(),
			                                             DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE,
			                                             IID_PPV_ARGS(&sharedB))))
			{
				sharedB->QueryInterface(IID_PPV_ARGS(&km));
				D3D11_TEXTURE2D_DESC sd{};
				sharedB->GetDesc(&sd);
				sd.Usage = D3D11_USAGE_STAGING;
				sd.BindFlags = 0;
				sd.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
				sd.MiscFlags = 0;
				ID3D11Texture2D* staging = nullptr;
				check(SUCCEEDED(devB->CreateTexture2D(&sd, nullptr, &staging)), "create staging texture");
				if (km != nullptr && staging != nullptr && km->AcquireSync(kVideoMutexKey, 1000) == S_OK)
				{
					ctxB->CopyResource(staging, sharedB);
					km->ReleaseSync(kVideoMutexKey);
					D3D11_MAPPED_SUBRESOURCE mapped{};
					if (SUCCEEDED(ctxB->Map(staging, 0, D3D11_MAP_READ, 0, &mapped)))
					{
						const auto* px = static_cast<const std::uint8_t*>(mapped.pData);
						std::printf("readback pixel0 = {%u,%u,%u,%u}\n", px[0], px[1], px[2], px[3]);
						check(near_byte(px[0], 51) && near_byte(px[1], 102) && near_byte(px[2], 153),
						      "shared texture carries the D3D12-rendered color");
						ctxB->Unmap(staging, 0);
					}
					else
					{
						check(false, "map staging texture");
					}
				}
				else
				{
					check(false, "acquire keyed mutex + copy shared texture");
				}
				release(staging);
			}
			else
			{
				check(false, "open shared texture by name");
			}
			release(km);
			release(sharedB);
			release(dev1);
		}
		release(ctxB);
		release(devB);

		// --- Performance regression guard: the D3D11On12 bridge + copy the hook does inside Present
		//     must stay off the present thread (measure a full frame with the hook live vs. removed). ---
		auto frame = [&] {
			const UINT idx = sc3->GetCurrentBackBufferIndex();
			allocator->Reset();
			cmdlist->Reset(allocator, nullptr);
			D3D12_RESOURCE_BARRIER b{};
			b.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
			b.Transition.pResource = render_targets[idx];
			b.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
			b.Transition.StateBefore = D3D12_RESOURCE_STATE_PRESENT;
			b.Transition.StateAfter = D3D12_RESOURCE_STATE_RENDER_TARGET;
			cmdlist->ResourceBarrier(1, &b);
			D3D12_CPU_DESCRIPTOR_HANDLE rtv = rtv_heap->GetCPUDescriptorHandleForHeapStart();
			rtv.ptr += static_cast<SIZE_T>(idx) * rtv_size;
			cmdlist->ClearRenderTargetView(rtv, kClear, 0, nullptr);
			std::swap(b.Transition.StateBefore, b.Transition.StateAfter);
			cmdlist->ResourceBarrier(1, &b);
			cmdlist->Close();
			ID3D12CommandList* lists[] = {cmdlist};
			queue->ExecuteCommandLists(1, lists);
		};
		auto present = [&] {
			swapchain->Present(0, 0);
			queue->Signal(fence, ++fence_value);
			if (fence->GetCompletedValue() < fence_value)
			{
				fence->SetEventOnCompletion(fence_value, fence_event);
				WaitForSingleObject(fence_event, 1000);
			}
		};
		const double hooked = cooptest::avg_present_ms(60, frame, present);
		hook::remove_present_hooks(); // baseline: same swapchain, hook removed
		const double base = cooptest::avg_present_ms(60, frame, present);
		std::printf("present-thread: hooked %.3f ms, unhooked %.3f ms, capture overhead %.3f ms\n", hooked, base,
		            hooked - base);
		check(hooked - base < cooptest::kPresentOverheadBudgetMs,
		      "D3D12 present hook stays off the present thread (overhead < one 60 Hz frame)");
	}

	hook::remove_present_hooks();

	if (fence_event != nullptr)
	{
		CloseHandle(fence_event);
	}
	release(fence);
	release(cmdlist);
	release(allocator);
	for (UINT i = 0; i < kBuffers; ++i)
	{
		release(render_targets[i]);
	}
	release(rtv_heap);
	release(sc3);
	release(swapchain);
	release(sc1);
	release(factory);
	release(queue);
	release(device);
	DestroyWindow(hwnd);
	UnregisterClassW(wc.lpszClassName, wc.hInstance);

	std::printf(g_failures == 0 ? "DX12 PRESENT HOOK TEST PASS\n" : "DX12 PRESENT HOOK TEST FAILED (%d)\n",
	            g_failures);
	return g_failures == 0 ? 0 : 1;
}