// Vulkan backend for the mock game. Brings up a real Vulkan instance/device/swap chain via
// volk (which dlopens vulkan-1.dll -- the loader-bypass case the capture hook must handle) and
// each frame clears the swap-chain image to the frame-counter colour with vkCmdClearColorImage
// (no pipeline, no shaders, no SPIR-V) and presents. The whole image encodes the frame number,
// so it animates and a dropped/stale frame is detectable by the capture test. Clear-only keeps
// this to one image-clear per frame; richer per-rect drawing would need a render pass.
#include "render_backend.hpp"

#include <vector>

#include <volk.h>

namespace coop::mock
{
namespace
{
class VkBackend : public RenderBackend
{
public:
	bool init(HWND hwnd, std::uint32_t width, std::uint32_t height) override
	{
		width_ = width;
		height_ = height;
		if (volkInitialize() != VK_SUCCESS)
		{
			return false; // no Vulkan loader on this machine
		}

		VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO};
		app.pApplicationName = "coop_mock_game";
		app.apiVersion = VK_API_VERSION_1_1;
		const char* inst_ext[] = {VK_KHR_SURFACE_EXTENSION_NAME, VK_KHR_WIN32_SURFACE_EXTENSION_NAME};
		VkInstanceCreateInfo ici{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO};
		ici.pApplicationInfo = &app;
		ici.enabledExtensionCount = 2;
		ici.ppEnabledExtensionNames = inst_ext;
		if (vkCreateInstance(&ici, nullptr, &instance_) != VK_SUCCESS)
		{
			return false;
		}
		volkLoadInstance(instance_);

		VkWin32SurfaceCreateInfoKHR sci{VK_STRUCTURE_TYPE_WIN32_SURFACE_CREATE_INFO_KHR};
		sci.hinstance = GetModuleHandleW(nullptr);
		sci.hwnd = hwnd;
		if (vkCreateWin32SurfaceKHR(instance_, &sci, nullptr, &surface_) != VK_SUCCESS)
		{
			return false;
		}

		if (!pick_device() || !create_device() || !create_swapchain() || !create_commands())
		{
			return false;
		}
		return true;
	}

	void render_and_present(std::uint32_t frame) override
	{
		if (device_ == VK_NULL_HANDLE || swapchain_ == VK_NULL_HANDLE)
		{
			return;
		}
		vkWaitForFences(device_, 1, &in_flight_, VK_TRUE, UINT64_MAX);

		std::uint32_t idx = 0;
		VkResult acq = vkAcquireNextImageKHR(device_, swapchain_, UINT64_MAX, acquire_sem_, VK_NULL_HANDLE, &idx);
		if (acq == VK_ERROR_OUT_OF_DATE_KHR || acq == VK_SUBOPTIMAL_KHR)
		{
			return; // skip this frame (the mock window isn't resized in practice)
		}
		if (acq != VK_SUCCESS)
		{
			return;
		}
		vkResetFences(device_, 1, &in_flight_);

		VkCommandBuffer cb = cmd_;
		vkResetCommandBuffer(cb, 0);
		VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO};
		bi.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
		vkBeginCommandBuffer(cb, &bi);

		barrier(cb, images_[idx], VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 0,
		        VK_ACCESS_TRANSFER_WRITE_BIT, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);

		std::uint8_t r = 0, g = 0, b = 0;
		frame_to_rgb(frame, r, g, b); // whole image encodes the frame -> animates + decodable
		VkClearColorValue cc{};
		cc.float32[0] = r / 255.0f;
		cc.float32[1] = g / 255.0f;
		cc.float32[2] = b / 255.0f;
		cc.float32[3] = 1.0f;
		VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};
		vkCmdClearColorImage(cb, images_[idx], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, &cc, 1, &range);

		barrier(cb, images_[idx], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
		        VK_ACCESS_TRANSFER_WRITE_BIT, 0, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT);
		vkEndCommandBuffer(cb);

		VkPipelineStageFlags wait_stage = VK_PIPELINE_STAGE_TRANSFER_BIT;
		VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO};
		si.waitSemaphoreCount = 1;
		si.pWaitSemaphores = &acquire_sem_;
		si.pWaitDstStageMask = &wait_stage;
		si.commandBufferCount = 1;
		si.pCommandBuffers = &cb;
		si.signalSemaphoreCount = 1;
		si.pSignalSemaphores = &submit_sem_;
		vkQueueSubmit(queue_, 1, &si, in_flight_);

		VkPresentInfoKHR pi{VK_STRUCTURE_TYPE_PRESENT_INFO_KHR};
		pi.waitSemaphoreCount = 1;
		pi.pWaitSemaphores = &submit_sem_;
		pi.swapchainCount = 1;
		pi.pSwapchains = &swapchain_;
		pi.pImageIndices = &idx;
		vkQueuePresentKHR(queue_, &pi);
	}

	[[nodiscard]] const char* name() const override
	{
		return "vk";
	}

	~VkBackend() override
	{
		if (device_ != VK_NULL_HANDLE)
		{
			vkDeviceWaitIdle(device_);
			if (in_flight_ != VK_NULL_HANDLE)
				vkDestroyFence(device_, in_flight_, nullptr);
			if (acquire_sem_ != VK_NULL_HANDLE)
				vkDestroySemaphore(device_, acquire_sem_, nullptr);
			if (submit_sem_ != VK_NULL_HANDLE)
				vkDestroySemaphore(device_, submit_sem_, nullptr);
			if (pool_ != VK_NULL_HANDLE)
				vkDestroyCommandPool(device_, pool_, nullptr);
			if (swapchain_ != VK_NULL_HANDLE)
				vkDestroySwapchainKHR(device_, swapchain_, nullptr);
			vkDestroyDevice(device_, nullptr);
		}
		if (surface_ != VK_NULL_HANDLE)
			vkDestroySurfaceKHR(instance_, surface_, nullptr);
		if (instance_ != VK_NULL_HANDLE)
			vkDestroyInstance(instance_, nullptr);
	}

private:
	bool pick_device()
	{
		std::uint32_t n = 0;
		vkEnumeratePhysicalDevices(instance_, &n, nullptr);
		std::vector<VkPhysicalDevice> devs(n);
		vkEnumeratePhysicalDevices(instance_, &n, devs.data());
		for (VkPhysicalDevice pd : devs)
		{
			std::uint32_t qn = 0;
			vkGetPhysicalDeviceQueueFamilyProperties(pd, &qn, nullptr);
			std::vector<VkQueueFamilyProperties> qf(qn);
			vkGetPhysicalDeviceQueueFamilyProperties(pd, &qn, qf.data());
			for (std::uint32_t i = 0; i < qn; ++i)
			{
				VkBool32 present = VK_FALSE;
				vkGetPhysicalDeviceSurfaceSupportKHR(pd, i, surface_, &present);
				if ((qf[i].queueFlags & VK_QUEUE_GRAPHICS_BIT) && present)
				{
					phys_ = pd;
					qfam_ = i;
					return true;
				}
			}
		}
		return false;
	}

	bool create_device()
	{
		float prio = 1.0f;
		VkDeviceQueueCreateInfo qci{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO};
		qci.queueFamilyIndex = qfam_;
		qci.queueCount = 1;
		qci.pQueuePriorities = &prio;
		const char* dev_ext[] = {VK_KHR_SWAPCHAIN_EXTENSION_NAME};
		VkDeviceCreateInfo dci{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO};
		dci.queueCreateInfoCount = 1;
		dci.pQueueCreateInfos = &qci;
		dci.enabledExtensionCount = 1;
		dci.ppEnabledExtensionNames = dev_ext;
		if (vkCreateDevice(phys_, &dci, nullptr, &device_) != VK_SUCCESS)
		{
			return false;
		}
		volkLoadDevice(device_);
		vkGetDeviceQueue(device_, qfam_, 0, &queue_);
		return true;
	}

	bool create_swapchain()
	{
		VkSurfaceCapabilitiesKHR caps{};
		vkGetPhysicalDeviceSurfaceCapabilitiesKHR(phys_, surface_, &caps);

		std::uint32_t fn = 0;
		vkGetPhysicalDeviceSurfaceFormatsKHR(phys_, surface_, &fn, nullptr);
		std::vector<VkSurfaceFormatKHR> formats(fn);
		vkGetPhysicalDeviceSurfaceFormatsKHR(phys_, surface_, &fn, formats.data());
		VkSurfaceFormatKHR chosen = formats.empty() ? VkSurfaceFormatKHR{VK_FORMAT_B8G8R8A8_UNORM,
		                                                                 VK_COLOR_SPACE_SRGB_NONLINEAR_KHR}
		                                             : formats[0];
		for (const VkSurfaceFormatKHR& f : formats)
		{
			if (f.format == VK_FORMAT_B8G8R8A8_UNORM || f.format == VK_FORMAT_R8G8B8A8_UNORM)
			{
				chosen = f;
				break;
			}
		}
		format_ = chosen.format;

		std::uint32_t want = caps.minImageCount + 1;
		if (caps.maxImageCount > 0 && want > caps.maxImageCount)
		{
			want = caps.maxImageCount;
		}
		VkSwapchainCreateInfoKHR sc{VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR};
		sc.surface = surface_;
		sc.minImageCount = want;
		sc.imageFormat = chosen.format;
		sc.imageColorSpace = chosen.colorSpace;
		sc.imageExtent = caps.currentExtent.width != 0xFFFFFFFFu
		                     ? caps.currentExtent
		                     : VkExtent2D{width_, height_};
		sc.imageArrayLayers = 1;
		// TRANSFER_DST so we can clear it; TRANSFER_SRC so the capture hook can copy it out.
		sc.imageUsage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
		sc.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
		sc.preTransform = caps.currentTransform;
		sc.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
		sc.presentMode = VK_PRESENT_MODE_FIFO_KHR; // vsync, universally supported
		sc.clipped = VK_TRUE;
		if (vkCreateSwapchainKHR(device_, &sc, nullptr, &swapchain_) != VK_SUCCESS)
		{
			return false;
		}
		std::uint32_t in = 0;
		vkGetSwapchainImagesKHR(device_, swapchain_, &in, nullptr);
		images_.resize(in);
		vkGetSwapchainImagesKHR(device_, swapchain_, &in, images_.data());
		return true;
	}

	bool create_commands()
	{
		VkCommandPoolCreateInfo pci{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO};
		pci.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
		pci.queueFamilyIndex = qfam_;
		if (vkCreateCommandPool(device_, &pci, nullptr, &pool_) != VK_SUCCESS)
		{
			return false;
		}
		VkCommandBufferAllocateInfo ai{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO};
		ai.commandPool = pool_;
		ai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
		ai.commandBufferCount = 1;
		if (vkAllocateCommandBuffers(device_, &ai, &cmd_) != VK_SUCCESS)
		{
			return false;
		}
		VkSemaphoreCreateInfo si{VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO};
		VkFenceCreateInfo fi{VK_STRUCTURE_TYPE_FENCE_CREATE_INFO};
		fi.flags = VK_FENCE_CREATE_SIGNALED_BIT;
		return vkCreateSemaphore(device_, &si, nullptr, &acquire_sem_) == VK_SUCCESS &&
		       vkCreateSemaphore(device_, &si, nullptr, &submit_sem_) == VK_SUCCESS &&
		       vkCreateFence(device_, &fi, nullptr, &in_flight_) == VK_SUCCESS;
	}

	static void barrier(VkCommandBuffer cb, VkImage img, VkImageLayout from, VkImageLayout to,
	                    VkAccessFlags src_access, VkAccessFlags dst_access, VkPipelineStageFlags src_stage,
	                    VkPipelineStageFlags dst_stage)
	{
		VkImageMemoryBarrier b{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER};
		b.srcAccessMask = src_access;
		b.dstAccessMask = dst_access;
		b.oldLayout = from;
		b.newLayout = to;
		b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
		b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
		b.image = img;
		b.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};
		vkCmdPipelineBarrier(cb, src_stage, dst_stage, 0, 0, nullptr, 0, nullptr, 1, &b);
	}

	std::uint32_t width_ = 0;
	std::uint32_t height_ = 0;
	VkInstance instance_ = VK_NULL_HANDLE;
	VkSurfaceKHR surface_ = VK_NULL_HANDLE;
	VkPhysicalDevice phys_ = VK_NULL_HANDLE;
	std::uint32_t qfam_ = 0;
	VkDevice device_ = VK_NULL_HANDLE;
	VkQueue queue_ = VK_NULL_HANDLE;
	VkSwapchainKHR swapchain_ = VK_NULL_HANDLE;
	VkFormat format_ = VK_FORMAT_UNDEFINED;
	std::vector<VkImage> images_;
	VkCommandPool pool_ = VK_NULL_HANDLE;
	VkCommandBuffer cmd_ = VK_NULL_HANDLE;
	VkSemaphore acquire_sem_ = VK_NULL_HANDLE;
	VkSemaphore submit_sem_ = VK_NULL_HANDLE;
	VkFence in_flight_ = VK_NULL_HANDLE;
};
} // namespace

std::unique_ptr<RenderBackend> create_vk_backend()
{
	return std::make_unique<VkBackend>();
}

} // namespace coop::mock