// CoopAllTheThings Vulkan capture layer.
//
// A real Vulkan *implicit layer* the loader inserts at vkCreateInstance -- guaranteed to be in
// the chain before the game resolves vkQueuePresentKHR. This is the reliable early-presence path
// for games that initialize Vulkan immediately (which inject-after-launch + the inline-hook
// vk_hook can't catch). The actual read-back (copy the presented image into the shared keyed-mutex
// texture, off the present thread) is the shared coop::hook::VkCapture; this file only gets the
// layer into the dispatch chain and feeds VkCapture each present.
//
// Scoping: an implicit layer loads into *every* Vulkan app, so the layer only *captures* when it
// recognises the process as the host's target -- env COOP_VK_LAYER_FORCE=1 (tests), or this
// process's image basename matches %TEMP%\coop_vk_target.txt (written by the host's "Set up Vulkan
// layer" checkbox). Otherwise it's a pure pass-through.
//
// The loader/layer interface structs (VkLayer*CreateInfo, VkNegotiateLayerInterface) live in
// vk_layer.h, which Vulkan-Headers doesn't ship, so they're declared here to the stable
// loader-interface-version-2 ABI.
#include <algorithm>
#include <atomic>
#include <cstdarg>
#include <cstdio>
#include <cstring>
#include <mutex>
#include <string>
#include <vector>

#include <windows.h>

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

#define VK_NO_PROTOTYPES
#define VK_USE_PLATFORM_WIN32_KHR
#include <vulkan/vulkan.h>

#include "coop/shared_memory.hpp"
#include "ipc_client.hpp"
#include "vk_capture.hpp"

// --- Loader/layer interface (interface version 2) ---------------------------
extern "C" {
typedef enum VkLayerFunction_ {
	COOP_VK_LAYER_LINK_INFO = 0,
	COOP_VK_LOADER_DATA_CALLBACK = 1,
	COOP_VK_LOADER_LAYER_CREATE_DEVICE_CALLBACK = 2,
	COOP_VK_LOADER_FEATURES = 3,
} CoopVkLayerFunction;

typedef PFN_vkVoidFunction(VKAPI_PTR* PFN_GetPhysicalDeviceProcAddr)(VkInstance, const char*);

typedef struct VkLayerInstanceLink_ {
	struct VkLayerInstanceLink_* pNext;
	PFN_vkGetInstanceProcAddr pfnNextGetInstanceProcAddr;
	PFN_GetPhysicalDeviceProcAddr pfnNextGetPhysicalDeviceProcAddr;
} VkLayerInstanceLink;

typedef struct VkLayerInstanceCreateInfo {
	VkStructureType sType; // 1000000000 = VK_STRUCTURE_TYPE_LOADER_INSTANCE_CREATE_INFO
	const void* pNext;
	CoopVkLayerFunction function;
	union {
		VkLayerInstanceLink* pLayerInfo;
		void* pfnCallback; // other callbacks (unused here); keeps the union pointer-sized
	} u;
} VkLayerInstanceCreateInfo;

typedef struct VkLayerDeviceLink_ {
	struct VkLayerDeviceLink_* pNext;
	PFN_vkGetInstanceProcAddr pfnNextGetInstanceProcAddr;
	PFN_vkGetDeviceProcAddr pfnNextGetDeviceProcAddr;
} VkLayerDeviceLink;

typedef struct VkLayerDeviceCreateInfo {
	VkStructureType sType; // 1000000001 = VK_STRUCTURE_TYPE_LOADER_DEVICE_CREATE_INFO
	const void* pNext;
	CoopVkLayerFunction function;
	union {
		VkLayerDeviceLink* pLayerInfo;
		void* pfnCallback;
	} u;
} VkLayerDeviceCreateInfo;

typedef struct VkNegotiateLayerInterface {
	uint32_t sType; // 1 = LAYER_NEGOTIATE_INTERFACE_STRUCT
	void* pNext;
	uint32_t loaderLayerInterfaceVersion;
	PFN_vkGetInstanceProcAddr pfnGetInstanceProcAddr;
	PFN_vkGetDeviceProcAddr pfnGetDeviceProcAddr;
	PFN_GetPhysicalDeviceProcAddr pfnGetPhysicalDeviceProcAddr;
} VkNegotiateLayerInterface;
}

namespace {
// The loader tags its chain-link structs with small, loader-internal sType values (not the
// 1000000000-range): VK_STRUCTURE_TYPE_LOADER_INSTANCE_CREATE_INFO = 47, _DEVICE = 48. These are
// from the (unvendored) vk_layer.h and are stable across loader versions.
constexpr VkStructureType kLoaderInstanceCreateInfo = static_cast<VkStructureType>(47);
constexpr VkStructureType kLoaderDeviceCreateInfo = static_cast<VkStructureType>(48);

using coop::hook::IpcClient;
using coop::hook::VkCapture;

IpcClient g_ipc;
std::atomic<bool> g_ipc_tried{false}; // reaper-thread side: connect once on the first published frame
bool g_active = false;                // do we capture in this process? (scoping)

// Chain dispatch.
PFN_vkGetInstanceProcAddr g_next_gipa = nullptr;
PFN_vkGetDeviceProcAddr g_next_gdpa = nullptr;
PFN_vkQueuePresentKHR g_real_present = nullptr;
PFN_vkCreateSwapchainKHR g_real_create_swapchain = nullptr;
PFN_vkGetSwapchainImagesKHR g_get_swapchain_images = nullptr; // for our CreateSwapchain image tracking

VkInstance g_instance = VK_NULL_HANDLE;
VkPhysicalDevice g_phys = VK_NULL_HANDLE;
VkDevice g_device = VK_NULL_HANDLE;
std::uint32_t g_qfam = 0;

VkCapture g_cap; // the shared, off-present-thread read-back

struct SwapInfo {
	VkSwapchainKHR sc;
	VkFormat fmt;
	std::uint32_t w, h;
	std::vector<VkImage> images;
};
std::vector<SwapInfo> g_swaps;
// Serialize all g_swaps access: create (push) and present (find) can run on different game threads,
// and DestroyDevice clears it. We don't track swapchain destruction (forwarding a destroy wrong could
// break the game); create de-dups by handle and an LRU cap bounds growth so repeated resize/recreate
// can't grow g_swaps unbounded or match a recycled handle's stale images.
std::mutex g_swaps_mutex;
constexpr std::size_t kMaxTrackedSwaps = 8;

bool eq(const char* a, const char* b)
{
	return std::strcmp(a, b) == 0;
}

// Optional file trace for debugging the chain dispatch (enable with COOP_VK_LAYER_LOG).
void logvk(const char* fmt, ...)
{
	static int enabled = -1;
	if (enabled < 0) {
		enabled = GetEnvironmentVariableW(L"COOP_VK_LAYER_LOG", nullptr, 0) != 0 ? 1 : 0;
	}
	if (enabled == 0) {
		return;
	}
	wchar_t dir[MAX_PATH] = {};
	if (GetTempPathW(MAX_PATH, dir) == 0) {
		return;
	}
	FILE* f = nullptr;
	if (_wfopen_s(&f, (std::wstring(dir) + L"coop_vk_layer.log").c_str(), L"a") != 0 || f == nullptr) {
		return;
	}
	va_list ap;
	va_start(ap, fmt);
	std::vfprintf(f, fmt, ap);
	va_end(ap);
	std::fputc('\n', f);
	std::fclose(f);
}

// Decide whether this process is the host's capture target (see file header).
bool decide_active()
{
	if (GetEnvironmentVariableW(L"COOP_VK_LAYER_FORCE", nullptr, 0) != 0) {
		return true;
	}
	wchar_t dir[MAX_PATH] = {};
	const DWORD n = GetTempPathW(MAX_PATH, dir);
	if (n == 0 || n >= MAX_PATH) {
		return false;
	}
	HANDLE f = CreateFileW((std::wstring(dir) + L"coop_vk_target.txt").c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr,
	                       OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
	if (f == INVALID_HANDLE_VALUE) {
		return false;
	}
	char want[MAX_PATH] = {};
	DWORD got = 0;
	ReadFile(f, want, sizeof(want) - 1, &got, nullptr);
	CloseHandle(f);
	// Trim trailing whitespace/newline.
	while (got > 0 && (want[got - 1] == '\n' || want[got - 1] == '\r' || want[got - 1] == ' ')) {
		want[--got] = '\0';
	}
	if (got == 0) {
		return false;
	}
	wchar_t self[MAX_PATH] = {};
	GetModuleFileNameW(nullptr, self, MAX_PATH);
	const wchar_t* base = wcsrchr(self, L'\\');
	base = base != nullptr ? base + 1 : self;
	char self8[MAX_PATH] = {};
	WideCharToMultiByte(CP_UTF8, 0, base, -1, self8, sizeof(self8), nullptr, nullptr);
	return _stricmp(self8, want) == 0;
}

// Caller must hold g_swaps_mutex; copy out what you need before unlocking (a concurrent push_back
// can reallocate and dangle the returned pointer).
const SwapInfo* find_swap(VkSwapchainKHR sc)
{
	for (const SwapInfo& s : g_swaps) {
		if (s.sc == sc) {
			return &s;
		}
	}
	return nullptr;
}

// Per-second present-rate trace (enable with COOP_VK_LAYER_LOG): confirms the game keeps its frame
// rate while capturing (the whole point of the off-present-thread read-back) and how many of those
// presents the reaper actually mirrored.
void trace_present_rate()
{
	static int enabled = -1;
	if (enabled < 0) {
		enabled = GetEnvironmentVariableW(L"COOP_VK_LAYER_LOG", nullptr, 0) != 0 ? 1 : 0;
	}
	if (enabled == 0) {
		return;
	}
	static LARGE_INTEGER freq{};
	static LARGE_INTEGER last{};
	static int count = 0;
	static std::uint64_t last_published = 0;
	if (freq.QuadPart == 0) {
		QueryPerformanceFrequency(&freq);
		QueryPerformanceCounter(&last);
	}
	++count;
	LARGE_INTEGER now;
	QueryPerformanceCounter(&now);
	const double sec = static_cast<double>(now.QuadPart - last.QuadPart) / static_cast<double>(freq.QuadPart);
	if (sec >= 1.0) {
		const std::uint64_t pub = g_cap.frames_published();
		logvk("present rate %.1f/s  captured %.1f/s (game keeps its rate; capture is off the present thread)",
		      count / sec, (pub - last_published) / sec);
		last = now;
		count = 0;
		last_published = pub;
	}
}

VKAPI_ATTR VkResult VKAPI_CALL layer_QueuePresentKHR(VkQueue queue, const VkPresentInfoKHR* pi)
{
	trace_present_rate(); // unconditional so an inactive (pass-through) run gives a baseline to compare
	if (g_active && g_ipc.connected()) {
		g_ipc.note_present();
	}
	if (g_active && pi != nullptr && pi->swapchainCount == 1) {
		// Copy the matched swapchain out under the lock, then capture without holding it.
		VkImage image = VK_NULL_HANDLE;
		VkFormat fmt = VK_FORMAT_UNDEFINED;
		std::uint32_t w = 0, h = 0;
		bool matched = false;
		{
			std::scoped_lock lock(g_swaps_mutex);
			const SwapInfo* s = find_swap(pi->pSwapchains[0]);
			const std::uint32_t idx = pi->pImageIndices[0];
			if (s != nullptr && idx < s->images.size()) {
				image = s->images[idx];
				fmt = s->fmt;
				w = s->w;
				h = s->h;
				matched = true;
			}
		}
		if (matched) {
			VkSemaphore chained = VK_NULL_HANDLE;
			if (g_cap.present(image, fmt, w, h, pi->pWaitSemaphores, pi->waitSemaphoreCount, chained)) {
				VkPresentInfoKHR p = *pi;
				p.waitSemaphoreCount = 1;
				p.pWaitSemaphores = &chained;
				return g_real_present(queue, &p);
			}
		}
	}
	return g_real_present(queue, pi);
}

VKAPI_ATTR VkResult VKAPI_CALL layer_CreateSwapchainKHR(VkDevice device, const VkSwapchainCreateInfoKHR* ci,
                                                        const VkAllocationCallbacks* a, VkSwapchainKHR* out)
{
	// Log the game's chosen present mode (= its sync mode) -- 0 IMMEDIATE, 1 MAILBOX, 2 FIFO (vsync),
	// 3 FIFO_RELAXED. We pass `ci` straight through, so whatever the game asked for is what it gets;
	// this proves the layer never changes vsync.
	logvk("CreateSwapchain: presentMode=%d (0=IMMEDIATE 1=MAILBOX 2=FIFO 3=FIFO_RELAXED) %ux%u minImageCount=%u",
	      static_cast<int>(ci->presentMode), ci->imageExtent.width, ci->imageExtent.height, ci->minImageCount);
	const VkResult r = g_real_create_swapchain(device, ci, a, out);
	if (g_active && r == VK_SUCCESS && out && g_get_swapchain_images) {
		SwapInfo info{};
		info.sc = *out;
		info.fmt = ci->imageFormat;
		info.w = ci->imageExtent.width;
		info.h = ci->imageExtent.height;
		std::uint32_t n = 0;
		g_get_swapchain_images(device, *out, &n, nullptr);
		info.images.resize(n);
		g_get_swapchain_images(device, *out, &n, info.images.data());
		{
			std::scoped_lock lock(g_swaps_mutex);
			g_swaps.erase(
			    std::remove_if(g_swaps.begin(), g_swaps.end(), [&](const SwapInfo& e) { return e.sc == info.sc; }),
			    g_swaps.end());
			g_swaps.push_back(std::move(info));
			if (g_swaps.size() > kMaxTrackedSwaps) {
				g_swaps.erase(g_swaps.begin());
			}
		}
	}
	return r;
}

// Resolve the device functions the read-back needs and start the capture component.
void start_capture(VkDevice dev)
{
	VkCapture::Fns f{};
#define LOAD(field, vkname) f.field = reinterpret_cast<PFN_##vkname>(g_next_gdpa(dev, #vkname))
	LOAD(GetDeviceQueue, vkGetDeviceQueue);
	LOAD(CreateCommandPool, vkCreateCommandPool);
	LOAD(DestroyCommandPool, vkDestroyCommandPool);
	LOAD(AllocateCommandBuffers, vkAllocateCommandBuffers);
	LOAD(BeginCommandBuffer, vkBeginCommandBuffer);
	LOAD(EndCommandBuffer, vkEndCommandBuffer);
	LOAD(ResetCommandBuffer, vkResetCommandBuffer);
	LOAD(CmdPipelineBarrier, vkCmdPipelineBarrier);
	LOAD(CmdCopyImageToBuffer, vkCmdCopyImageToBuffer);
	LOAD(QueueSubmit, vkQueueSubmit);
	LOAD(CreateFence, vkCreateFence);
	LOAD(DestroyFence, vkDestroyFence);
	LOAD(WaitForFences, vkWaitForFences);
	LOAD(ResetFences, vkResetFences);
	LOAD(GetFenceStatus, vkGetFenceStatus);
	LOAD(CreateSemaphore, vkCreateSemaphore);
	LOAD(DestroySemaphore, vkDestroySemaphore);
	LOAD(CreateBuffer, vkCreateBuffer);
	LOAD(DestroyBuffer, vkDestroyBuffer);
	LOAD(GetBufferMemoryRequirements, vkGetBufferMemoryRequirements);
	LOAD(AllocateMemory, vkAllocateMemory);
	LOAD(FreeMemory, vkFreeMemory);
	LOAD(BindBufferMemory, vkBindBufferMemory);
	LOAD(MapMemory, vkMapMemory);
	LOAD(UnmapMemory, vkUnmapMemory);
	LOAD(InvalidateMappedMemoryRanges, vkInvalidateMappedMemoryRanges);
	LOAD(DeviceWaitIdle, vkDeviceWaitIdle);
#undef LOAD
	f.GetPhysicalDeviceMemoryProperties = reinterpret_cast<PFN_vkGetPhysicalDeviceMemoryProperties>(
	    g_next_gipa(g_instance, "vkGetPhysicalDeviceMemoryProperties"));
	g_get_swapchain_images = reinterpret_cast<PFN_vkGetSwapchainImagesKHR>(g_next_gdpa(dev, "vkGetSwapchainImagesKHR"));

	g_cap.init(g_phys, dev, g_qfam, f, GetCurrentProcessId(), [](std::uint32_t w, std::uint32_t h) {
		// Runs on the reaper thread after each frame is published to the shared texture. Connect the
		// IPC channel lazily here (the host may not have created it yet at device-create time), then
		// publish the frame so the host's generation counter advances.
		if (!g_ipc_tried.exchange(true)) {
			g_ipc.connect(/*attempts=*/40, /*delay_ms=*/25);
		}
		if (g_ipc.connected()) {
			g_ipc.publish_video_frame(w, h, static_cast<std::uint32_t>(DXGI_FORMAT_R8G8B8A8_UNORM));
		}
	});
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL layer_gdpa(VkDevice device, const char* name);

VKAPI_ATTR VkResult VKAPI_CALL layer_CreateDevice(VkPhysicalDevice phys, const VkDeviceCreateInfo* ci,
                                                  const VkAllocationCallbacks* a, VkDevice* out)
{
	auto* link = reinterpret_cast<VkLayerDeviceCreateInfo*>(const_cast<void*>(ci->pNext));
	while (link != nullptr && !(link->sType == kLoaderDeviceCreateInfo && link->function == COOP_VK_LAYER_LINK_INFO)) {
		link = reinterpret_cast<VkLayerDeviceCreateInfo*>(const_cast<void*>(link->pNext));
	}
	if (link == nullptr) {
		logvk("CreateDevice: LINK_INFO not found");
		return VK_ERROR_INITIALIZATION_FAILED;
	}
	PFN_vkGetInstanceProcAddr next_gipa = link->u.pLayerInfo->pfnNextGetInstanceProcAddr;
	PFN_vkGetDeviceProcAddr next_gdpa = link->u.pLayerInfo->pfnNextGetDeviceProcAddr;
	link->u.pLayerInfo = link->u.pLayerInfo->pNext; // advance the chain for the next layer
	auto create = reinterpret_cast<PFN_vkCreateDevice>(next_gipa(g_instance, "vkCreateDevice"));
	const VkResult r = create(phys, ci, a, out);
	logvk("CreateDevice: result=%d active=%d", (int)r, g_active ? 1 : 0);
	if (r == VK_SUCCESS && out != nullptr && g_device == VK_NULL_HANDLE) {
		g_phys = phys;
		g_device = *out;
		g_next_gdpa = next_gdpa;
		g_qfam = ci->queueCreateInfoCount > 0 ? ci->pQueueCreateInfos[0].queueFamilyIndex : 0;
		g_real_present = reinterpret_cast<PFN_vkQueuePresentKHR>(next_gdpa(*out, "vkQueuePresentKHR"));
		g_real_create_swapchain = reinterpret_cast<PFN_vkCreateSwapchainKHR>(next_gdpa(*out, "vkCreateSwapchainKHR"));
		if (g_active) {
			start_capture(*out);
		}
	}
	return r;
}

VKAPI_ATTR VkResult VKAPI_CALL layer_CreateInstance(const VkInstanceCreateInfo* ci, const VkAllocationCallbacks* a,
                                                    VkInstance* out)
{
	auto* link = reinterpret_cast<VkLayerInstanceCreateInfo*>(const_cast<void*>(ci->pNext));
	while (link != nullptr
	       && !(link->sType == kLoaderInstanceCreateInfo && link->function == COOP_VK_LAYER_LINK_INFO)) {
		link = reinterpret_cast<VkLayerInstanceCreateInfo*>(const_cast<void*>(link->pNext));
	}
	if (link == nullptr) {
		logvk("CreateInstance: LINK_INFO not found");
		return VK_ERROR_INITIALIZATION_FAILED;
	}
	PFN_vkGetInstanceProcAddr next_gipa = link->u.pLayerInfo->pfnNextGetInstanceProcAddr;
	link->u.pLayerInfo = link->u.pLayerInfo->pNext; // advance the chain
	auto create = reinterpret_cast<PFN_vkCreateInstance>(next_gipa(nullptr, "vkCreateInstance"));
	const VkResult r = create(ci, a, out);
	if (r == VK_SUCCESS && out != nullptr) {
		g_instance = *out;
		g_next_gipa = next_gipa;
		g_active = decide_active();
	}
	logvk("CreateInstance: result=%d active=%d", (int)r, g_active ? 1 : 0);
	return r;
}

VKAPI_ATTR void VKAPI_CALL layer_DestroyDevice(VkDevice device, const VkAllocationCallbacks* a)
{
	auto destroy = reinterpret_cast<PFN_vkDestroyDevice>(g_next_gdpa(device, "vkDestroyDevice"));
	if (g_active && device == g_device) {
		g_cap.shutdown(); // joins the reaper, drains the device, frees the read-back resources
		{
			std::scoped_lock lock(g_swaps_mutex);
			g_swaps.clear();
		}
		g_device = VK_NULL_HANDLE;
	}
	destroy(device, a);
}

VKAPI_ATTR void VKAPI_CALL layer_DestroyInstance(VkInstance instance, const VkAllocationCallbacks* a)
{
	auto destroy = reinterpret_cast<PFN_vkDestroyInstance>(g_next_gipa(instance, "vkDestroyInstance"));
	g_instance = VK_NULL_HANDLE;
	destroy(instance, a);
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL layer_gdpa(VkDevice device, const char* name)
{
	if (name == nullptr) {
		return nullptr;
	}
	if (eq(name, "vkGetDeviceProcAddr"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_gdpa);
	if (eq(name, "vkQueuePresentKHR"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_QueuePresentKHR);
	if (eq(name, "vkCreateSwapchainKHR"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_CreateSwapchainKHR);
	if (eq(name, "vkDestroyDevice"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_DestroyDevice);
	return g_next_gdpa != nullptr ? g_next_gdpa(device, name) : nullptr;
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL layer_gipa(VkInstance instance, const char* name)
{
	if (name == nullptr) {
		return nullptr;
	}
	if (eq(name, "vkGetInstanceProcAddr"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_gipa);
	if (eq(name, "vkCreateInstance"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_CreateInstance);
	if (eq(name, "vkCreateDevice"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_CreateDevice);
	if (eq(name, "vkDestroyInstance"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_DestroyInstance);
	if (eq(name, "vkGetDeviceProcAddr"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_gdpa);
	return g_next_gipa != nullptr ? g_next_gipa(instance, name) : nullptr;
}

} // namespace

extern "C" __declspec(dllexport) VkResult VKAPI_CALL
vkNegotiateLoaderLayerInterfaceVersion(VkNegotiateLayerInterface* pVersionStruct)
{
	logvk("negotiate: requestedVersion=%u", pVersionStruct->loaderLayerInterfaceVersion);
	if (pVersionStruct->loaderLayerInterfaceVersion > 2) {
		pVersionStruct->loaderLayerInterfaceVersion = 2;
	}
	pVersionStruct->pfnGetInstanceProcAddr = layer_gipa;
	pVersionStruct->pfnGetDeviceProcAddr = layer_gdpa;
	pVersionStruct->pfnGetPhysicalDeviceProcAddr = nullptr;
	return VK_SUCCESS;
}

// Also export the entry points directly, for loaders that probe them by name.
extern "C" __declspec(dllexport) PFN_vkVoidFunction VKAPI_CALL coop_vkGetInstanceProcAddr(VkInstance i, const char* n)
{
	return layer_gipa(i, n);
}
extern "C" __declspec(dllexport) PFN_vkVoidFunction VKAPI_CALL coop_vkGetDeviceProcAddr(VkDevice d, const char* n)
{
	return layer_gdpa(d, n);
}