// coop_vk_validate: validate the injected Vulkan capture backend against a REAL game.
//
// Vulkan can't be late-hooked (it caches its present pointer at init), so capture needs early
// presence. This drives both productized early-presence methods end-to-end against a shipping
// Vulkan title (Sphere Spectacle, Steam appid 1123040 by default):
//   layer  - register the implicit coop_vk_layer (scoped to the game image), launch via Steam, and
//            read the frames it publishes. The realistic path for a Steam-launched game.
//   inject - suspended-launch the game's exe directly, inject coop_hook.dll, resume (the hook arms
//            before vkCreateInstance). Only works if the title runs when launched outside Steam.
//
// For each method it asserts capture works (frames advance), the image is sane (non-black, the
// captured resolution), saves a BMP screenshot for visual confirmation, and measures the game's
// present rate while capturing (a healthy rate at the refresh cap = capture stays off the critical
// path). Usage: coop_vk_validate <layer|inject> [seconds] [exe-path]
#include <cstdint>
#include <cstdio>
#include <string>
#include <vector>

#include <windows.h>

#include <shellapi.h>
#include <tlhelp32.h>

#include <d3d11.h>

#include "capture/shared_texture.hpp"
#include "coop/protocol.hpp"
#include "coop/shared_memory.hpp"
#include "coop/tool_paths.hpp"

using namespace coop;

namespace
{
constexpr const wchar_t* kDefaultExe =
    L"G:\\SteamLibrary\\steamapps\\common\\Sphere Spectacle\\sphere.exe";
constexpr const wchar_t* kSteamUrl = L"steam://rungameid/1123040";

unsigned long find_pid(const wchar_t* image)
{
	unsigned long pid = 0;
	HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (snap == INVALID_HANDLE_VALUE)
	{
		return 0;
	}
	PROCESSENTRY32W pe{};
	pe.dwSize = sizeof(pe);
	for (BOOL ok = Process32FirstW(snap, &pe); ok; ok = Process32NextW(snap, &pe))
	{
		if (_wcsicmp(pe.szExeFile, image) == 0)
		{
			pid = pe.th32ProcessID;
			break;
		}
	}
	CloseHandle(snap);
	return pid;
}

void kill_pid(unsigned long pid)
{
	if (HANDLE h = OpenProcess(PROCESS_TERMINATE, FALSE, pid))
	{
		TerminateProcess(h, 0);
		CloseHandle(h);
	}
}

// Register / unregister the implicit Vulkan layer, scoped to `image_basename`. Mirrors
// host vk_layer_setup but resolves the manifest from the deployable root (one dir up from tools/).
void layer_register(const std::wstring& image_basename)
{
	wchar_t tmp[MAX_PATH] = {};
	if (GetTempPathW(MAX_PATH, tmp) != 0)
	{
		const std::wstring sf = std::wstring(tmp) + L"coop_vk_target.txt";
		char utf8[260] = {};
		const int n = WideCharToMultiByte(CP_UTF8, 0, image_basename.c_str(), -1, utf8, sizeof(utf8), nullptr,
		                                  nullptr);
		HANDLE f = CreateFileW(sf.c_str(), GENERIC_WRITE, 0, nullptr, CREATE_ALWAYS, 0, nullptr);
		if (f != INVALID_HANDLE_VALUE)
		{
			DWORD wr = 0;
			WriteFile(f, utf8, n > 0 ? static_cast<DWORD>(n - 1) : 0, &wr, nullptr); // drop the NUL
			CloseHandle(f);
		}
	}
	HKEY key = nullptr;
	if (RegCreateKeyExW(HKEY_CURRENT_USER, L"SOFTWARE\\Khronos\\Vulkan\\ImplicitLayers", 0, nullptr, 0,
	                    KEY_SET_VALUE, nullptr, &key, nullptr) == ERROR_SUCCESS)
	{
		const std::wstring mp = deployed_artifact_path(L"coop_vk_layer.json");
		DWORD enabled = 0;
		RegSetValueExW(key, mp.c_str(), 0, REG_DWORD, reinterpret_cast<const BYTE*>(&enabled), sizeof(enabled));
		RegCloseKey(key);
	}
}

void layer_unregister()
{
	HKEY key = nullptr;
	if (RegOpenKeyExW(HKEY_CURRENT_USER, L"SOFTWARE\\Khronos\\Vulkan\\ImplicitLayers", 0, KEY_SET_VALUE, &key) ==
	    ERROR_SUCCESS)
	{
		RegDeleteValueW(key, deployed_artifact_path(L"coop_vk_layer.json").c_str());
		RegCloseKey(key);
	}
	wchar_t tmp[MAX_PATH] = {};
	if (GetTempPathW(MAX_PATH, tmp) != 0)
	{
		DeleteFileW((std::wstring(tmp) + L"coop_vk_target.txt").c_str());
	}
}

bool inject(unsigned long pid)
{
	const std::wstring dll = deployed_artifact_path(L"coop_hook.dll");
	HANDLE process = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | PROCESS_VM_OPERATION |
	                                 PROCESS_VM_WRITE | PROCESS_VM_READ,
	                             FALSE, pid);
	if (process == nullptr)
	{
		return false;
	}
	const SIZE_T bytes = (dll.size() + 1) * sizeof(wchar_t);
	void* remote = VirtualAllocEx(process, nullptr, bytes, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
	bool ok = false;
	if (remote != nullptr && WriteProcessMemory(process, remote, dll.c_str(), bytes, nullptr))
	{
		auto load = reinterpret_cast<LPTHREAD_START_ROUTINE>(
		    GetProcAddress(GetModuleHandleW(L"kernel32.dll"), "LoadLibraryW"));
		if (HANDLE th = CreateRemoteThread(process, nullptr, 0, load, remote, 0, nullptr))
		{
			WaitForSingleObject(th, INFINITE);
			DWORD code = 0;
			GetExitCodeThread(th, &code);
			CloseHandle(th);
			ok = code != 0;
		}
	}
	if (remote != nullptr)
	{
		VirtualFreeEx(process, remote, 0, MEM_RELEASE);
	}
	CloseHandle(process);
	return ok;
}

SharedBlock* make_ipc(SharedMemory& shm, unsigned long pid)
{
	if (!shm.create(shared_memory_name(pid), sizeof(SharedBlock)))
	{
		return nullptr;
	}
	auto* b = shm.as<SharedBlock>();
	b->version = kProtocolVersion;
	b->pad_count = 0;
	b->sequence.store(0, std::memory_order_relaxed);
	for (std::uint32_t s = 0; s < HookSubsys_Count; ++s)
	{
		b->control.subsystem_disabled[s].store(0, std::memory_order_release); // all on (video included)
	}
	b->magic = kProtocolMagic;
	return b;
}

ID3D11Device* make_device()
{
	ID3D11Device* dev = nullptr;
	const D3D_FEATURE_LEVEL fl[] = {D3D_FEATURE_LEVEL_11_1, D3D_FEATURE_LEVEL_11_0};
	if (FAILED(D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, 0, fl,
	                             static_cast<UINT>(std::size(fl)), D3D11_SDK_VERSION, &dev, nullptr, nullptr)))
	{
		return nullptr;
	}
	return dev;
}

bool write_bmp(const std::wstring& path, const std::vector<std::uint8_t>& rgba, std::uint32_t w, std::uint32_t h)
{
	const std::uint32_t row = (w * 3 + 3) & ~3u;
	const std::uint32_t imgsize = row * h;
	BITMAPFILEHEADER fh{};
	BITMAPINFOHEADER ih{};
	fh.bfType = 0x4D42;
	fh.bfOffBits = sizeof(fh) + sizeof(ih);
	fh.bfSize = fh.bfOffBits + imgsize;
	ih.biSize = sizeof(ih);
	ih.biWidth = static_cast<LONG>(w);
	ih.biHeight = static_cast<LONG>(h);
	ih.biPlanes = 1;
	ih.biBitCount = 24;
	ih.biCompression = BI_RGB;
	ih.biSizeImage = imgsize;
	HANDLE f = CreateFileW(path.c_str(), GENERIC_WRITE, 0, nullptr, CREATE_ALWAYS, 0, nullptr);
	if (f == INVALID_HANDLE_VALUE)
	{
		return false;
	}
	DWORD wr = 0;
	WriteFile(f, &fh, sizeof(fh), &wr, nullptr);
	WriteFile(f, &ih, sizeof(ih), &wr, nullptr);
	std::vector<std::uint8_t> line(row, 0);
	for (int y = static_cast<int>(h) - 1; y >= 0; --y) // BMP is bottom-up
	{
		for (std::uint32_t x = 0; x < w; ++x)
		{
			const std::uint8_t* p = &rgba[(static_cast<std::size_t>(y) * w + x) * 4];
			line[x * 3 + 0] = p[2]; // B
			line[x * 3 + 1] = p[1]; // G
			line[x * 3 + 2] = p[0]; // R
		}
		WriteFile(f, line.data(), row, &wr, nullptr);
	}
	CloseHandle(f);
	return true;
}

VideoShareView read_share(const SharedBlock* b)
{
	VideoShareView v;
	v.generation = b->video.generation.load(std::memory_order_acquire);
	v.width = b->video.width;
	v.height = b->video.height;
	v.format = b->video.format;
	v.present_calls = b->video.present_calls;
	return v;
}
} // namespace

int main(int argc, char** argv)
{
	const std::string mode = argc > 1 ? argv[1] : "layer";
	const int seconds = argc > 2 ? std::atoi(argv[2]) : 12;
	std::wstring exe = kDefaultExe;
	if (argc > 3)
	{
		const std::string a = argv[3];
		exe.assign(a.begin(), a.end());
	}
	const bool layer_mode = mode != "inject";
	std::printf("== Vulkan backend validation: method=%s game=%ls ==\n", layer_mode ? "layer" : "inject",
	            exe.c_str());

	int failures = 0;
	auto check = [&](bool ok, const char* what) {
		std::printf("%s %s\n", ok ? "  ok:" : "FAIL:", what);
		if (!ok)
		{
			++failures;
		}
	};

	// Clean slate.
	if (unsigned long old = find_pid(L"sphere.exe"))
	{
		kill_pid(old);
		Sleep(1000);
	}

	PROCESS_INFORMATION pi{};
	unsigned long pid = 0;
	if (layer_mode)
	{
		layer_register(L"sphere.exe");
		ShellExecuteW(nullptr, L"open", kSteamUrl, nullptr, nullptr, SW_SHOWNORMAL);
		std::printf("   launched via Steam; waiting for sphere.exe...\n");
		for (int i = 0; i < 40 && pid == 0; ++i)
		{
			Sleep(500);
			pid = find_pid(L"sphere.exe");
		}
	}
	else
	{
		STARTUPINFOW si{};
		si.cb = sizeof(si);
		std::wstring cmd = exe;
		// Launch with the game's own folder as the working directory -- the game loads steam_api64.dll
		// and resources/ relative to cwd, so launching with our cwd makes it fail to initialize (and
		// produce no presents), which earlier looked like "requires Steam". It does not.
		const std::size_t slash = exe.find_last_of(L"\\/");
		const std::wstring workdir = slash != std::wstring::npos ? exe.substr(0, slash) : std::wstring();
		if (!CreateProcessW(exe.c_str(), cmd.data(), nullptr, nullptr, FALSE, CREATE_SUSPENDED, nullptr,
		                    workdir.empty() ? nullptr : workdir.c_str(), &si, &pi))
		{
			check(false, "suspended-launch the game exe directly");
			return 1;
		}
		pid = pi.dwProcessId;
	}

	if (pid == 0)
	{
		check(false, "game process appeared");
		layer_unregister();
		return 1;
	}
	std::printf("   game pid=%lu\n", pid);

	// Create the IPC block right away so the layer/hook can connect + publish present counts.
	SharedMemory shm;
	SharedBlock* block = make_ipc(shm, pid);
	if (block == nullptr)
	{
		check(false, "create IPC block");
		kill_pid(pid);
		layer_unregister();
		return 1;
	}

	if (!layer_mode)
	{
		const bool injected = inject(pid);
		ResumeThread(pi.hThread);
		check(injected, "inject coop_hook.dll early (pre-vkCreateInstance)");
		CloseHandle(pi.hThread);
	}
	else
	{
		// The layer does video; also inject coop_hook.dll (late) so its FOCUS subsystem keeps the
		// game rendering at full rate while unfocused -- the realistic co-injected setup, and
		// required for a meaningful present-rate measurement (an unfocused game throttles itself,
		// which is exactly why the hook spoofs focus). Its Vulkan video hook is inert (too late);
		// the layer is the video producer.
		Sleep(2000); // let the game create its window first
		const bool injected = inject(pid);
		std::printf("   focus/input hook injected (late, for focus spoof): %d\n", injected ? 1 : 0);
		Sleep(1500); // let the hook attach + spoof focus before we measure
	}

	ID3D11Device* device = make_device();
	if (device == nullptr)
	{
		check(false, "create a D3D11 device to read the shared texture");
		kill_pid(pid);
		layer_unregister();
		return 1;
	}

	SharedTextureSource src;
	src.init(device);

	// Capture loop: wait for frames, track advance + non-black, grab a screenshot.
	std::uint32_t first_gen = 0, last_gen = 0, cap_w = 0, cap_h = 0;
	std::uint64_t nonblack_frames = 0;
	std::vector<std::uint8_t> shot;
	std::uint32_t shot_w = 0, shot_h = 0;
	const DWORD end = GetTickCount() + static_cast<DWORD>(seconds) * 1000;
	bool alive = true;
	while (GetTickCount() < end && (alive = find_pid(L"sphere.exe") == pid))
	{
		Sleep(50);
		const VideoShareView sv = read_share(block);
		if (!src.update(sv, pid))
		{
			continue;
		}
		cap_w = src.width();
		cap_h = src.height();
		if (first_gen == 0)
		{
			first_gen = sv.generation;
		}
		last_gen = sv.generation;
		// Sample a few pixels for non-black; grab a full screenshot mid-run.
		std::uint8_t px[4] = {};
		if (src.read_pixel(cap_w / 2, cap_h / 2, px) && (px[0] | px[1] | px[2]) != 0)
		{
			++nonblack_frames;
		}
		if (shot.empty() && src.frames_copied() > 10)
		{
			src.read_frame(shot, shot_w, shot_h);
		}
	}

	const std::uint64_t presents = block->video.present_calls;
	std::printf("   captured %ux%u  copied=%llu  present_calls=%llu  gen %u..%u  nonblack=%llu\n", cap_w, cap_h,
	            static_cast<unsigned long long>(src.frames_copied()),
	            static_cast<unsigned long long>(presents), first_gen, last_gen,
	            static_cast<unsigned long long>(nonblack_frames));

	// Not-applicable skip: some titles produce no Vulkan presents when their exe is suspended-launched
	// directly (e.g. they refuse to run without a real Steam launch). For those the inject method
	// can't reach the game and the layer method should be used. (Sphere Spectacle does NOT need this:
	// it runs fine launched directly with its own folder as the working directory -- see the inject
	// launch above -- so this branch should not trigger for it.)
	if (!layer_mode && presents == 0 && src.frames_copied() == 0)
	{
		std::printf("   the directly-launched exe produced no Vulkan presents -- it may require a real Steam\n"
		            "   launch (or failed to initialize), so the suspended-inject method can't reach it here;\n"
		            "   use the layer method. (The early-inject mechanism itself is covered by mock_game_test.)\n");
		kill_pid(pid);
		device->Release();
		std::printf("SKIP vk_validate (inject produced no presents for this title)\n");
		return 0;
	}

	check(alive, "game stayed alive through capture");
	check(presents > 0, "hook/layer saw the game's Vulkan presents");
	check(src.frames_copied() >= 10, "host copied many shared frames (capture works)");
	check(last_gen > first_gen + 5, "captured frames advance (live mirror, not a stuck frame)");
	check(cap_w >= 320 && cap_h >= 240, "captured a sensible resolution");
	check(nonblack_frames >= 5, "captured frames are non-black (real image content)");

	// Screenshot for visual confirmation of correctness (colors / brightness / no swizzle).
	if (!shot.empty())
	{
		const std::wstring out = exe_directory() + (layer_mode ? L"vk_validate_layer.bmp" : L"vk_validate_inject.bmp");
		if (write_bmp(out, shot, shot_w, shot_h))
		{
			std::printf("   screenshot: %ls (%ux%u)\n", out.c_str(), shot_w, shot_h);
		}
	}
	else
	{
		check(false, "grabbed a screenshot frame");
	}

	// Performance gate (this is the check the earlier version refused to make -- it reported the rate
	// and rationalized it, which hid the 144->3 FPS stall). The capture runs off the present thread,
	// so the game must keep a healthy present rate while we mirror. The capture is NOT throttled --
	// it follows the present rate, which vsync paces -- so the mirror rate tracks the present rate. A
	// present rate that collapses (the bug was ~3/s) fails here.
	if (alive && find_pid(L"sphere.exe") == pid)
	{
		const std::uint64_t p0 = block->video.present_calls;
		const std::uint32_t g0 = block->video.generation.load(std::memory_order_acquire);
		Sleep(3000);
		const std::uint64_t p1 = block->video.present_calls;
		const std::uint32_t g1 = block->video.generation.load(std::memory_order_acquire);
		const double fps = static_cast<double>(p1 - p0) / 3.0;
		// The hook bumps video.generation on every mirrored frame (independent of the host reading),
		// so its delta is the true mirror rate even while we're just sleeping here.
		const double mirror = static_cast<double>(g1 - g0) / 3.0;
		std::printf("   present rate while capturing = %.1f /s; mirror rate = %.1f /s (capture is off the "
		            "present thread; the mirror follows the present rate -- vsync paces it)\n",
		            fps, mirror);
		check(fps > 30.0, "game keeps a healthy present rate while capturing (no present-thread stall)");
	}

	kill_pid(pid);
	device->Release();
	if (layer_mode)
	{
		layer_unregister();
	}

	std::printf(failures == 0 ? "PASS vk_validate (%s)\n" : "FAILED vk_validate (%s, %d)\n",
	            layer_mode ? "layer" : "inject", failures);
	return failures == 0 ? 0 : 1;
}