#include <Cplt/Entrypoint/Backend.hpp>

#if BUILD_CORE_WITH_VULKAN_BACKEND
#	include <iostream>
#	include <stdexcept>

#	define GLFW_INCLUDE_NONE
#	define GLFW_INCLUDE_VULKAN
#	include <GLFW/glfw3.h>

#	include <backend/imgui_impl_glfw.h>
#	include <backend/imgui_impl_vulkan.h>
#	include <backend/imgui_impl_vulkan.cpp>

class VulkanBackend : public RenderingBackend
{
private:
	GLFWwindow* mWindow;

	VkAllocationCallbacks* mAllocator = NULL;
	VkInstance mInstance = VK_NULL_HANDLE;
	VkPhysicalDevice mPhysicalDevice = VK_NULL_HANDLE;
	VkDevice mDevice = VK_NULL_HANDLE;
	uint32_t mQueueFamily = (uint32_t)-1;
	VkQueue mQueue = VK_NULL_HANDLE;
	VkDebugReportCallbackEXT mDebugReport = VK_NULL_HANDLE;
	VkPipelineCache mPipelineCache = VK_NULL_HANDLE;
	VkDescriptorPool mDescriptorPool = VK_NULL_HANDLE;

	ImGui_ImplVulkanH_Window mMainWindowData;
	int mMinImageCount = 2;
	bool mSwapChainRebuild = false;

public:
	VulkanBackend()
	{
		glfwSetErrorCallback(&GlfwErrorCallback);
		if (!glfwInit()) {
			throw std::runtime_error("Failed to initialize GLFW.");
		}

		glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
		mWindow = glfwCreateWindow(1280, 720, "Cplt", nullptr, nullptr);
		if (mWindow == nullptr) {
			throw std::runtime_error("Failed to create GLFW window.");
		}

		if (!glfwVulkanSupported()) {
			throw std::runtime_error("GLFW reports vulkan not supported.");
		}

		uint32_t extensionsCount = 0;
		const char** extensions = glfwGetRequiredInstanceExtensions(&extensionsCount);
		SetupVulkan(extensions, extensionsCount);

		// Create window surface
		VkSurfaceKHR surface;
		VkResult err = glfwCreateWindowSurface(mInstance, mWindow, mAllocator, &surface);
		CheckVkResults(err);

		// Create framebuffers
		int w, h;
		glfwGetFramebufferSize(mWindow, &w, &h);
		SetupVulkanWindow(&mMainWindowData, surface, w, h);

		IMGUI_CHECKVERSION();
		ImGui::CreateContext();

		ImGui_ImplGlfw_InitForVulkan(mWindow, true);
		ImGui_ImplVulkan_InitInfo init_info = {};
		init_info.Instance = mInstance;
		init_info.PhysicalDevice = mPhysicalDevice;
		init_info.Device = mDevice;
		init_info.QueueFamily = mQueueFamily;
		init_info.Queue = mQueue;
		init_info.PipelineCache = mPipelineCache;
		init_info.DescriptorPool = mDescriptorPool;
		init_info.Allocator = mAllocator;
		init_info.MinImageCount = mMinImageCount;
		init_info.ImageCount = mMainWindowData.ImageCount;
		init_info.CheckVkResultFn = CheckVkResults;
		ImGui_ImplVulkan_Init(&init_info, mMainWindowData.RenderPass);
	}

	virtual ~VulkanBackend()
	{
		auto err = vkDeviceWaitIdle(mDevice);
		CheckVkResults(err);
		ImGui_ImplVulkan_Shutdown();
		ImGui_ImplGlfw_Shutdown();
		ImGui::DestroyContext();

		CleanupVulkanWindow();
		CleanupVulkan();

		glfwDestroyWindow(mWindow);
		glfwTerminate();
	}

	virtual void RunUntilWindowClose(void (*windowContent)()) override
	{
		// Upload Fonts
		{
			// Use any command queue
			VkCommandPool commandPool = mMainWindowData.Frames[mMainWindowData.FrameIndex].CommandPool;
			VkCommandBuffer commandBuffer = mMainWindowData.Frames[mMainWindowData.FrameIndex].CommandBuffer;

			CheckVkResults(vkResetCommandPool(mDevice, commandPool, 0));
			VkCommandBufferBeginInfo beginInfo = {};
			beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
			beginInfo.flags |= VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
			CheckVkResults(vkBeginCommandBuffer(commandBuffer, &beginInfo));

			ImGui_ImplVulkan_CreateFontsTexture(commandBuffer);

			VkSubmitInfo endInfo = {};
			endInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
			endInfo.commandBufferCount = 1;
			endInfo.pCommandBuffers = &commandBuffer;
			CheckVkResults(vkEndCommandBuffer(commandBuffer));
			CheckVkResults(vkQueueSubmit(mQueue, 1, &endInfo, VK_NULL_HANDLE));

			CheckVkResults(vkDeviceWaitIdle(mDevice));
			ImGui_ImplVulkan_DestroyFontUploadObjects();
		}

		while (!glfwWindowShouldClose(mWindow)) {
			glfwPollEvents();

			// Resize swap chain?
			if (mSwapChainRebuild) {
				int width, height;
				glfwGetFramebufferSize(mWindow, &width, &height);
				if (width > 0 && height > 0) {
					ImGui_ImplVulkan_SetMinImageCount(mMinImageCount);
					ImGui_ImplVulkanH_CreateOrResizeWindow(mInstance, mPhysicalDevice, mDevice, &mMainWindowData, mQueueFamily, mAllocator, width, height, mMinImageCount);
					mMainWindowData.FrameIndex = 0;
					mSwapChainRebuild = false;
				}
			}

			// Start the Dear ImGui frame
			ImGui_ImplVulkan_NewFrame();
			ImGui_ImplGlfw_NewFrame();
			ImGui::NewFrame();

			windowContent();

			ImGui::Render();
			ImDrawData* drawData = ImGui::GetDrawData();
			const bool isMinimized = (drawData->DisplaySize.x <= 0.0f || drawData->DisplaySize.y <= 0.0f);
			if (!isMinimized) {
				const ImVec4 kClearColor = ImVec4(0.45f, 0.55f, 0.60f, 1.00f);
				mMainWindowData.ClearValue.color.float32[0] = kClearColor.x * kClearColor.w;
				mMainWindowData.ClearValue.color.float32[1] = kClearColor.y * kClearColor.w;
				mMainWindowData.ClearValue.color.float32[2] = kClearColor.z * kClearColor.w;
				mMainWindowData.ClearValue.color.float32[3] = kClearColor.w;
				FrameRender(&mMainWindowData, drawData);
				FramePresent(&mMainWindowData);
			}
		}
	}

private:
	void SetupVulkan(const char** extensions, uint32_t extensions_count)
	{
		VkResult err;

		// Create Vulkan Instance
		{
			VkInstanceCreateInfo createInfo = {};
			createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
			createInfo.enabledExtensionCount = extensions_count;
			createInfo.ppEnabledExtensionNames = extensions;
			// Create Vulkan Instance without any debug feature
			err = vkCreateInstance(&createInfo, mAllocator, &mInstance);
			CheckVkResults(err);
		}

		// Select GPU
		{
			uint32_t gpuCount;
			err = vkEnumeratePhysicalDevices(mInstance, &gpuCount, NULL);
			CheckVkResults(err);
			IM_ASSERT(gpuCount > 0);

			VkPhysicalDevice* gpus = (VkPhysicalDevice*)malloc(sizeof(VkPhysicalDevice) * gpuCount);
			err = vkEnumeratePhysicalDevices(mInstance, &gpuCount, gpus);
			CheckVkResults(err);

			// If a number >1 of GPUs got reported, find discrete GPU if present, or use first one available. This covers
			// most common cases (multi-gpu/integrated+dedicated graphics). Handling more complicated setups (multiple
			// dedicated GPUs) is out of scope of this sample.
			int useGpu = 0;
			for (int i = 0; i < (int)gpuCount; i++)
			{
				VkPhysicalDeviceProperties properties;
				vkGetPhysicalDeviceProperties(gpus[i], &properties);
				if (properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU)
				{
					useGpu = i;
					break;
				}
			}

			mPhysicalDevice = gpus[useGpu];
			free(gpus);
		}

		// Select graphics queue family
		{
			uint32_t count;
			vkGetPhysicalDeviceQueueFamilyProperties(mPhysicalDevice, &count, NULL);

			auto queues = std::make_unique<VkQueueFamilyProperties[]>(count);
			vkGetPhysicalDeviceQueueFamilyProperties(mPhysicalDevice, &count, queues.get());
			for (uint32_t i = 0; i < count; i++) {
				if (queues[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)
				{
					mQueueFamily = i;
					break;
				}
			}

			IM_ASSERT(mQueueFamily != (uint32_t)-1);
		}

		// Create Logical Device (with 1 queue)
		{
			int deviceExtensionCount = 1;
			const char* deviceExtensions[] = { "VK_KHR_swapchain" };
			const float queuePriority[] = { 1.0f };
			VkDeviceQueueCreateInfo queue_info[1] = {};
			queue_info[0].sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
			queue_info[0].queueFamilyIndex = mQueueFamily;
			queue_info[0].queueCount = 1;
			queue_info[0].pQueuePriorities = queuePriority;
			VkDeviceCreateInfo createInfo = {};
			createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
			createInfo.queueCreateInfoCount = sizeof(queue_info) / sizeof(queue_info[0]);
			createInfo.pQueueCreateInfos = queue_info;
			createInfo.enabledExtensionCount = deviceExtensionCount;
			createInfo.ppEnabledExtensionNames = deviceExtensions;
			err = vkCreateDevice(mPhysicalDevice, &createInfo, mAllocator, &mDevice);
			CheckVkResults(err);
			vkGetDeviceQueue(mDevice, mQueueFamily, 0, &mQueue);
		}

		// Create Descriptor Pool
		{
			VkDescriptorPoolSize poolSizes[] = {
				{ VK_DESCRIPTOR_TYPE_SAMPLER, 1000 },
				{ VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1000 },
				{ VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, 1000 },
				{ VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1000 },
				{ VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, 1000 },
				{ VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, 1000 },
				{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1000 },
				{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1000 },
				{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, 1000 },
				{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, 1000 },
				{ VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 1000 }
			};
			VkDescriptorPoolCreateInfo poolInfo = {};
			poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
			poolInfo.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
			poolInfo.maxSets = 1000 * IM_ARRAYSIZE(poolSizes);
			poolInfo.poolSizeCount = (uint32_t)IM_ARRAYSIZE(poolSizes);
			poolInfo.pPoolSizes = poolSizes;
			err = vkCreateDescriptorPool(mDevice, &poolInfo, mAllocator, &mDescriptorPool);
			CheckVkResults(err);
		}
	}

	void SetupVulkanWindow(ImGui_ImplVulkanH_Window* wd, VkSurfaceKHR surface, int width, int height)
	{
		wd->Surface = surface;

		// Check for WSI support
		VkBool32 res;
		vkGetPhysicalDeviceSurfaceSupportKHR(mPhysicalDevice, mQueueFamily, wd->Surface, &res);
		if (res != VK_TRUE) {
			throw "Error no WSI support on physical device 0.";
		}

		// Select Surface Format
		const VkFormat requestSurfaceImageFormat[] = { VK_FORMAT_B8G8R8A8_UNORM, VK_FORMAT_R8G8B8A8_UNORM, VK_FORMAT_B8G8R8_UNORM, VK_FORMAT_R8G8B8_UNORM };
		const VkColorSpaceKHR requestSurfaceColorSpace = VK_COLORSPACE_SRGB_NONLINEAR_KHR;
		wd->SurfaceFormat = ImGui_ImplVulkanH_SelectSurfaceFormat(mPhysicalDevice, wd->Surface, requestSurfaceImageFormat, (size_t)IM_ARRAYSIZE(requestSurfaceImageFormat), requestSurfaceColorSpace);

		// Select Present Mode
		VkPresentModeKHR present_modes[] = { VK_PRESENT_MODE_FIFO_KHR };
		wd->PresentMode = ImGui_ImplVulkanH_SelectPresentMode(mPhysicalDevice, wd->Surface, &present_modes[0], IM_ARRAYSIZE(present_modes));

		// Create SwapChain, RenderPass, Framebuffer, etc.
		IM_ASSERT(mMinImageCount >= 2);
		ImGui_ImplVulkanH_CreateOrResizeWindow(mInstance, mPhysicalDevice, mDevice, wd, mQueueFamily, mAllocator, width, height, mMinImageCount);
	}

	void FrameRender(ImGui_ImplVulkanH_Window* wd, ImDrawData* drawData)
	{
		VkResult err;

		VkSemaphore imageAcquiredSemaphore = wd->FrameSemaphores[wd->SemaphoreIndex].ImageAcquiredSemaphore;
		VkSemaphore renderCompleteSemaphore = wd->FrameSemaphores[wd->SemaphoreIndex].RenderCompleteSemaphore;
		err = vkAcquireNextImageKHR(mDevice, wd->Swapchain, UINT64_MAX, imageAcquiredSemaphore, VK_NULL_HANDLE, &wd->FrameIndex);
		if (err == VK_ERROR_OUT_OF_DATE_KHR || err == VK_SUBOPTIMAL_KHR) {
			mSwapChainRebuild = true;
			return;
		}
		CheckVkResults(err);

		ImGui_ImplVulkanH_Frame* fd = &wd->Frames[wd->FrameIndex];
		{
			err = vkWaitForFences(mDevice, 1, &fd->Fence, VK_TRUE, UINT64_MAX); // wait indefinitely instead of periodically checking
			CheckVkResults(err);

			err = vkResetFences(mDevice, 1, &fd->Fence);
			CheckVkResults(err);
		}
		{
			err = vkResetCommandPool(mDevice, fd->CommandPool, 0);
			CheckVkResults(err);
			VkCommandBufferBeginInfo info = {};
			info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
			info.flags |= VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
			err = vkBeginCommandBuffer(fd->CommandBuffer, &info);
			CheckVkResults(err);
		}
		{
			VkRenderPassBeginInfo info = {};
			info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
			info.renderPass = wd->RenderPass;
			info.framebuffer = fd->Framebuffer;
			info.renderArea.extent.width = wd->Width;
			info.renderArea.extent.height = wd->Height;
			info.clearValueCount = 1;
			info.pClearValues = &wd->ClearValue;
			vkCmdBeginRenderPass(fd->CommandBuffer, &info, VK_SUBPASS_CONTENTS_INLINE);
		}

		// Record dear imgui primitives into command buffer
		ImGui_ImplVulkan_RenderDrawData(drawData, fd->CommandBuffer);

		// Submit command buffer
		vkCmdEndRenderPass(fd->CommandBuffer);
		{
			VkPipelineStageFlags wait_stage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
			VkSubmitInfo info = {};
			info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
			info.waitSemaphoreCount = 1;
			info.pWaitSemaphores = &imageAcquiredSemaphore;
			info.pWaitDstStageMask = &wait_stage;
			info.commandBufferCount = 1;
			info.pCommandBuffers = &fd->CommandBuffer;
			info.signalSemaphoreCount = 1;
			info.pSignalSemaphores = &renderCompleteSemaphore;

			err = vkEndCommandBuffer(fd->CommandBuffer);
			CheckVkResults(err);
			err = vkQueueSubmit(mQueue, 1, &info, fd->Fence);
			CheckVkResults(err);
		}
	}

	void FramePresent(ImGui_ImplVulkanH_Window* wd)
	{
		if (mSwapChainRebuild) {
			return;
		}

		VkSemaphore renderCompleteSemaphore = wd->FrameSemaphores[wd->SemaphoreIndex].RenderCompleteSemaphore;
		VkPresentInfoKHR info = {};
		info.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
		info.waitSemaphoreCount = 1;
		info.pWaitSemaphores = &renderCompleteSemaphore;
		info.swapchainCount = 1;
		info.pSwapchains = &wd->Swapchain;
		info.pImageIndices = &wd->FrameIndex;
		VkResult err = vkQueuePresentKHR(mQueue, &info);
		if (err == VK_ERROR_OUT_OF_DATE_KHR || err == VK_SUBOPTIMAL_KHR) {
			mSwapChainRebuild = true;
			return;
		}
		CheckVkResults(err);
		wd->SemaphoreIndex = (wd->SemaphoreIndex + 1) % wd->ImageCount; // Now we can use the next set of semaphores
	}

	void CleanupVulkan()
	{
		vkDestroyDescriptorPool(mDevice, mDescriptorPool, mAllocator);

		vkDestroyDevice(mDevice, mAllocator);
		vkDestroyInstance(mInstance, mAllocator);
	}

	void CleanupVulkanWindow()
	{
		ImGui_ImplVulkanH_DestroyWindow(mInstance, mDevice, &mMainWindowData, mAllocator);
	}

	static void CheckVkResults(VkResult err)
	{
		if (err == 0) return;

		std::string message;
		message += "Vulkan error: VkResult = ";
		message += err;

		if (err < 0) {
			throw std::runtime_error(message);
		} else {
			std::cerr << message << '\n';
		}
	}
	static void GlfwErrorCallback(int errorCode, const char* message)
	{
		std::cerr << "GLFW Error " << errorCode << ": " << message << "\n";
	}
};

std::unique_ptr<RenderingBackend> RenderingBackend::CreateVulkanBackend()
{
	try {
		return std::make_unique<VulkanBackend>();
	} catch (std::exception& e) {
		return nullptr;
	}
}

#else // ^^ BUILD_CORE_WITH_VULKAN_BACKEND | ~BUILD_CORE_WITH_VULKAN_BACKEND vv

std::unique_ptr<RenderingBackend> RenderingBackend::CreateVulkanBackend()
{
	return nullptr;
}

#endif