#include "VulkanContext.h"
#include "VulkanCommandBuffers.h"
#include "VulkanDeviceManager.h"
#include "VulkanFramebuffers.h"
#include "VulkanSwapChain.h"
#include "VulkanVertexBuffer.h"
#include "imgui_impl_glfw.h"
#include "imgui_impl_vulkan.h"
#include "utilities/Logger.h"
#include <GLFW/glfw3.h>
#include <cstdint>
#include <vector>
#include <vulkan/vulkan_core.h>

VulkanContext::VulkanContext()
{
}

VulkanContext::~VulkanContext()
{
}

void VulkanContext::Initialize(FVulkanConfig& InConfig, const std::vector<Vertex>& InVertices)
{
	Config = InConfig;

	if (Config.bValidationEnabled)
	{
		InstanceManager.CreateInstance(&DebugManager);
		DebugManager.Initialize(InstanceManager.GetInstance());
	}
	else
	{
		InstanceManager.CreateInstance();
	}

	CreateSurface(InConfig.Window);

	DeviceManager.Initialize(FDeviceConfig(
		InstanceManager.GetInstance(),
		Config.bValidationEnabled,
		Surface,
		InConfig.Window));
	DeviceManager.PickPhysicalDevice();
	DeviceManager.CreateLogicalDevice();

	auto SwapChainSupport = DeviceManager.QuerySwapChainSupport(DeviceManager.GetPhysicalDevice());

	SwapChain.Initialize(FSwapConfig(
		DeviceManager.GetDevice(),
		Surface,
		Config.Window,
		DeviceManager.GetPhysicalQueueFamilies().GraphicsFamily,
		DeviceManager.GetPhysicalQueueFamilies().PresentFamily,
		SwapChainSupport.Capabilities,
		SwapChainSupport.Formats,
		SwapChainSupport.PresentModes));

	SwapChain.CreateSwapChain();
	SwapChain.CreateImageViews();

	RenderPass.Initialize(DeviceManager.GetDevice());
	RenderPass.CreateRenderPass(SwapChain.GetSwapChainImageFormat());

	GraphicsPipeline.Initialize(DeviceManager.GetDevice());
	GraphicsPipeline.CreateGraphicsPipeline(SwapChain.GetSwapChainExtent(), RenderPass.GetRenderPass());

	Framebuffers.Initialize(FFramebufferConfig(
		DeviceManager.GetDevice(),
		RenderPass.GetRenderPass(),
		SwapChain.GetSwapChainImageViews(),
		SwapChain.GetSwapChainExtent()));
	Framebuffers.CreateFramebuffers();

	CommandBuffers.Initialize(DeviceManager.GetDevice(), RenderPass.GetRenderPass());
	CommandBuffers.CreateCommandPool(DeviceManager.GetPhysicalQueueFamilies().GraphicsFamily);

	VertexBuffer.Initialize(FVertexBufferConfig(DeviceManager.GetDevice(), DeviceManager.GetPhysicalDevice()));
	VertexBuffer.CreateVertexBuffer(InVertices);

	CommandBuffers.CreateCommandBuffers(MAX_FRAMES_IN_FLIGHT);

	CreateSyncObjects();

	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;
	PoolInfo.poolSizeCount = (uint32_t)std::size(PoolSizes);
	PoolInfo.pPoolSizes = PoolSizes;
	vkCreateDescriptorPool(DeviceManager.GetDevice(), &PoolInfo, nullptr, &ImGuiPool);
}

void VulkanContext::Cleanup()
{
	if (bImGuiVulkanInitialized)
	{
		ImGui_ImplVulkan_Shutdown();
	}

	if (bImGuiGlfwInitialized)
	{
		ImGui_ImplGlfw_Shutdown();
	}

	CleanupSwapChain();

	VertexBuffer.Cleanup();

	GraphicsPipeline.Cleanup();
	RenderPass.Cleanup();

	for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
	{
		vkDestroySemaphore(DeviceManager.GetDevice(), ImageAvailableSemaphores[i], nullptr);
		vkDestroySemaphore(DeviceManager.GetDevice(), RenderFinishedSemaphores[i], nullptr);
		vkDestroyFence(DeviceManager.GetDevice(), InFlightFences[i], nullptr);
	}

	CommandBuffers.Cleanup();

	DeviceManager.Cleanup();

	if (Config.bValidationEnabled)
	{
		DebugManager.Cleanup();
	}

	vkDestroySurfaceKHR(InstanceManager.GetInstance(), Surface, nullptr);
	InstanceManager.Cleanup();
}

void VulkanContext::CreateSurface(GLFWwindow* Window)
{
	if (!Window)
	{
		Log::Error("Window not initialized.");
	}

	if (!InstanceManager.GetInstance())
	{
		Log::Error("Instance is null.");
	}

	VkResult result = glfwCreateWindowSurface(InstanceManager.GetInstance(), Window, nullptr, &Surface);
	if (result != VK_SUCCESS)
	{
		std::string errorMsg;
		switch (result)
		{
			case VK_ERROR_EXTENSION_NOT_PRESENT:
				errorMsg = "VK_ERROR_EXTENSION_NOT_PRESENT - Required extension not present";
				break;
			case VK_ERROR_INITIALIZATION_FAILED:
				errorMsg = "VK_ERROR_INITIALIZATION_FAILED - Initialization failed";
				break;
			case VK_ERROR_NATIVE_WINDOW_IN_USE_KHR:
				errorMsg = "VK_ERROR_NATIVE_WINDOW_IN_USE_KHR - Native window already in use";
				break;
			default:
				errorMsg = "Unknown error code: " + std::to_string(result);
				break;
		}

		Log::Error("Failed to create window surface: " + errorMsg);
	}
	else
	{
		Log::Info("Window surface created successfully.");
	}
}

void VulkanContext::CreateSyncObjects()
{
	ImageAvailableSemaphores.resize(MAX_FRAMES_IN_FLIGHT);
	RenderFinishedSemaphores.resize(MAX_FRAMES_IN_FLIGHT);
	InFlightFences.resize(MAX_FRAMES_IN_FLIGHT);

	VkSemaphoreCreateInfo SemaphoreInfo{};
	SemaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;

	VkFenceCreateInfo FenceInfo{};
	FenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
	FenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;

	for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
	{
		if (vkCreateSemaphore(DeviceManager.GetDevice(), &SemaphoreInfo, nullptr, &ImageAvailableSemaphores[i]) != VK_SUCCESS
			|| vkCreateSemaphore(DeviceManager.GetDevice(), &SemaphoreInfo, nullptr, &RenderFinishedSemaphores[i]) != VK_SUCCESS
			|| vkCreateFence(DeviceManager.GetDevice(), &FenceInfo, nullptr, &InFlightFences[i]) != VK_SUCCESS)
		{
			Log::Error("Failed to create semaphores!");
		}
		else
		{
			Log::Info("Successfully created semaphores");
		}
	}
}

void VulkanContext::DrawFrame(bool bDrawImGui, uint32_t InVerticesSize)
{
	vkWaitForFences(DeviceManager.GetDevice(), 1, &InFlightFences[CurrentFrame], VK_TRUE, UINT64_MAX);

	uint32_t ImageIndex;
	VkResult result = vkAcquireNextImageKHR(
		DeviceManager.GetDevice(),
		SwapChain.GetSwapChain(),
		UINT64_MAX,
		ImageAvailableSemaphores[CurrentFrame],
		VK_NULL_HANDLE,
		&ImageIndex);

	if (result == VK_ERROR_OUT_OF_DATE_KHR)
	{
		RecreateSwapChain();
		return;
	}
	else if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR)
	{
		Log::Error("Failed to acquire swap chain images!");
	}

	vkResetFences(DeviceManager.GetDevice(), 1, &InFlightFences[CurrentFrame]);

	vkResetCommandBuffer(CommandBuffers.GetCommandBuffer(CurrentFrame), 0);

	FRecordCommandBuffersParams Params{
		CommandBuffers.GetCommandBuffer(CurrentFrame),
		ImageIndex,
		VertexBuffer.GetVertexBuffer(),
		InVerticesSize,
		RenderPass.GetRenderPass(),
		SwapChain.GetSwapChainExtent(),
		GraphicsPipeline.GetGraphicsPipeline(),
		Framebuffers.GetSwapChainFrameBuffers(),
		bDrawImGui
		// DrawData
	};

	CommandBuffers.RecordCommandBuffer(Params);

	VkSubmitInfo SubmitInfo{};
	SubmitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;

	VkSemaphore			 WaitSemaphores[] = { ImageAvailableSemaphores[CurrentFrame] };
	VkPipelineStageFlags WaitStages[] = { VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT };
	SubmitInfo.waitSemaphoreCount = 1;
	SubmitInfo.pWaitSemaphores = WaitSemaphores;
	SubmitInfo.pWaitDstStageMask = WaitStages;

	SubmitInfo.commandBufferCount = 1;
	VkCommandBuffer CommandBuffer = CommandBuffers.GetCommandBuffer(CurrentFrame);
	SubmitInfo.pCommandBuffers = &CommandBuffer;

	VkSemaphore SignalSemaphores[] = { RenderFinishedSemaphores[CurrentFrame] };
	SubmitInfo.signalSemaphoreCount = 1;
	SubmitInfo.pSignalSemaphores = SignalSemaphores;

	if (vkQueueSubmit(DeviceManager.GetGraphicsQueue(), 1, &SubmitInfo, InFlightFences[CurrentFrame]) != VK_SUCCESS)
	{
		Log::Error("Failed to submit draw command buffer!");
	}

	VkPresentInfoKHR PresentInfo{};
	PresentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
	PresentInfo.waitSemaphoreCount = 1;
	PresentInfo.pWaitSemaphores = SignalSemaphores;

	VkSwapchainKHR SwapChains[] = { SwapChain.GetSwapChain() };

	PresentInfo.swapchainCount = 1;
	PresentInfo.pSwapchains = SwapChains;
	PresentInfo.pImageIndices = &ImageIndex;
	PresentInfo.pResults = nullptr;

	result = vkQueuePresentKHR(DeviceManager.GetPresentQueue(), &PresentInfo);

	if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || bFramebufferResized)
	{
		bFramebufferResized = false;
		RecreateSwapChain();
	}
	else if (result != VK_SUCCESS)
	{
		Log::Error("Failed to present swap chain image!");
	}

	CurrentFrame = (CurrentFrame + 1) % MAX_FRAMES_IN_FLIGHT;
}

void VulkanContext::RecreateSwapChain()
{
	int Width = 0, Height = 0;
	glfwGetFramebufferSize(Config.Window, &Width, &Height);
	while (Width == 0 || Height == 0)
	{
		glfwGetFramebufferSize(Config.Window, &Width, &Height);
		glfwWaitEvents();
	}

	Log::Info("Recreating SwapChain...");
	vkDeviceWaitIdle(DeviceManager.GetDevice());

	CleanupSwapChain();

	auto SwapChainSupport = DeviceManager.QuerySwapChainSupport(DeviceManager.GetPhysicalDevice());

	SwapChain.Initialize(FSwapConfig(
		DeviceManager.GetDevice(),
		Surface,
		Config.Window,
		DeviceManager.GetPhysicalQueueFamilies().GraphicsFamily,
		DeviceManager.GetPhysicalQueueFamilies().PresentFamily,
		SwapChainSupport.Capabilities,
		SwapChainSupport.Formats,
		SwapChainSupport.PresentModes));
	SwapChain.CreateSwapChain();
	SwapChain.CreateImageViews();

	Framebuffers.Initialize(FFramebufferConfig(
		DeviceManager.GetDevice(),
		RenderPass.GetRenderPass(),
		SwapChain.GetSwapChainImageViews(),
		SwapChain.GetSwapChainExtent()));
	Framebuffers.CreateFramebuffers();

	InitImGui();
}

void VulkanContext::InitImGui()
{
	if (!bImGuiGlfwInitialized)
	{
		if (!ImGui_ImplGlfw_InitForVulkan(Config.Window, true))
		{
			Log::Error("Failed to initialize ImGui GLFW backend!");
			return;
		}
		bImGuiGlfwInitialized = true;
	}

	if (bImGuiVulkanInitialized)
	{
		ImGui_ImplVulkan_Shutdown();
		bImGuiVulkanInitialized = false;
	}

	ImGui_ImplVulkan_InitInfo ImGuiInitInfo = {};
	ImGuiInitInfo.Instance = InstanceManager.GetInstance();
	ImGuiInitInfo.PhysicalDevice = DeviceManager.GetPhysicalDevice();
	ImGuiInitInfo.Device = DeviceManager.GetDevice();
	ImGuiInitInfo.QueueFamily = DeviceManager.GetPhysicalQueueFamilies().GraphicsFamily.value_or(0);
	ImGuiInitInfo.Queue = DeviceManager.GetGraphicsQueue();
	ImGuiInitInfo.PipelineCache = VK_NULL_HANDLE;
	ImGuiInitInfo.DescriptorPool = ImGuiPool;
	ImGuiInitInfo.MinImageCount = MAX_FRAMES_IN_FLIGHT;
	ImGuiInitInfo.ImageCount = SwapChain.GetSwapChainImageViews().size();
	ImGuiInitInfo.Allocator = nullptr;
	ImGuiInitInfo.PipelineInfoMain.RenderPass = RenderPass.GetRenderPass();
	ImGuiInitInfo.PipelineInfoMain.Subpass = 0;
	ImGuiInitInfo.PipelineInfoMain.MSAASamples = VK_SAMPLE_COUNT_1_BIT;
	ImGuiInitInfo.CheckVkResultFn = nullptr;

	if (!ImGui_ImplVulkan_Init(&ImGuiInitInfo))
	{
		Log::Error("Failed to reinitialize ImGui Vulkan backend!");
		return;
	}

	bImGuiVulkanInitialized = true;
}

void VulkanContext::CleanupSwapChain()
{
	Framebuffers.Cleanup();
	SwapChain.Cleanup();
}
