#include "VulkanSwapChain.h"

#include <vector>
#include <cstdint>	 // Necessary for uint32_t
#include <limits>	 // Necessary for std::numeric_limits
#include <algorithm> // Necessary for std::clamp
#include <vulkan/vulkan_core.h>

#include "utilities/Logger.h"

VulkanSwapChain::VulkanSwapChain()
{
}

VulkanSwapChain::~VulkanSwapChain()
{
	// Cleanup();
}

void VulkanSwapChain::Initialize(FSwapConfig InSwapConfig)
{
	SwapConfig = InSwapConfig;
}

void VulkanSwapChain::Cleanup()
{
	for (auto ImageView : SwapChainImageViews)
	{
		vkDestroyImageView(SwapConfig.Device, ImageView, nullptr);
	}
	vkDestroySwapchainKHR(SwapConfig.Device, SwapChain, nullptr);
}

VkSurfaceFormatKHR VulkanSwapChain::ChooseSwapSurfaceFormat(const std::vector<VkSurfaceFormatKHR>& AvailableFormats)
{
	for (const auto& AvailableFormat : AvailableFormats)
	{
		if (AvailableFormat.format == VK_FORMAT_B8G8R8A8_SRGB && AvailableFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR)
		{
			return AvailableFormat;
		}
	}
	return AvailableFormats[0];
}

VkPresentModeKHR VulkanSwapChain::ChooseSwapPresentMode(const std::vector<VkPresentModeKHR>& AvailablePresentModes)
{
	for (const auto& AvailablePresentMode : AvailablePresentModes)
	{
		if (AvailablePresentMode == VK_PRESENT_MODE_MAILBOX_KHR)
		{
			return AvailablePresentMode;
		}
	}
	return VK_PRESENT_MODE_FIFO_KHR;
}

VkExtent2D VulkanSwapChain::ChooseSwapExtent(const VkSurfaceCapabilitiesKHR& Capabilities)
{
	if (Capabilities.currentExtent.width != (std::numeric_limits<uint32_t>::max)())
	{
		return Capabilities.currentExtent;
	}
	else
	{
		if (SwapConfig.Window == nullptr)
		{
			Log::Error("GLFW window is null in CreateSwapChain!");
		}
		int Width, Height;
		glfwGetFramebufferSize(SwapConfig.Window, &Width, &Height);

		VkExtent2D ActualExtent = {
			static_cast<uint32_t>(Width),
			static_cast<uint32_t>(Height)
		};

		ActualExtent.width = std::clamp(ActualExtent.width, Capabilities.minImageExtent.width, Capabilities.maxImageExtent.width);
		ActualExtent.height = std::clamp(ActualExtent.height, Capabilities.minImageExtent.height, Capabilities.maxImageExtent.height);

		return ActualExtent;
	}
}

void VulkanSwapChain::CreateSwapChain()
{
	VkSurfaceFormatKHR SurfaceFormat = ChooseSwapSurfaceFormat(SwapConfig.Formats);
	VkPresentModeKHR   PresentMode = ChooseSwapPresentMode(SwapConfig.PresentModes);
	VkExtent2D		   Extent = ChooseSwapExtent(SwapConfig.Capabilities);

	uint32_t ImageCount = SwapConfig.Capabilities.minImageCount + 1;

	if (SwapConfig.Capabilities.maxImageCount > 0 && ImageCount > SwapConfig.Capabilities.maxImageCount)
	{
		ImageCount = SwapConfig.Capabilities.maxImageCount;
	}

	VkSwapchainCreateInfoKHR CreateInfo{};
	CreateInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
	CreateInfo.surface = SwapConfig.Surface;

	CreateInfo.minImageCount = ImageCount;
	CreateInfo.imageFormat = SurfaceFormat.format;
	CreateInfo.imageColorSpace = SurfaceFormat.colorSpace;
	CreateInfo.imageExtent = Extent;
	CreateInfo.imageArrayLayers = 1;
	CreateInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT; // may need VK_IMAGE_USAGE_TRANSFER_DST_BIT for post processing https://vulkan-tutorial.com/Drawing_a_triangle/Presentation/Swap_chain#:~:text=VK%5FIMAGE%5FUSAGE%5FTRANSFER%5FDST%5FBIT

	uint32_t QueueFamilyIndices[] = { SwapConfig.GraphicsFamily.value(),
		SwapConfig.PresentFamily.value() };

	if (SwapConfig.GraphicsFamily != SwapConfig.PresentFamily)
	{
		CreateInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
		CreateInfo.queueFamilyIndexCount = 2;
		CreateInfo.pQueueFamilyIndices = QueueFamilyIndices;
	}
	else
	{
		CreateInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
		CreateInfo.queueFamilyIndexCount = 0;
		CreateInfo.pQueueFamilyIndices = nullptr;
	}

	CreateInfo.preTransform = SwapConfig.Capabilities.currentTransform;
	CreateInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
	CreateInfo.presentMode = PresentMode;
	CreateInfo.clipped = VK_TRUE;
	CreateInfo.oldSwapchain = VK_NULL_HANDLE;

	if (vkCreateSwapchainKHR(SwapConfig.Device, &CreateInfo, nullptr, &SwapChain) != VK_SUCCESS)
	{
		Log::Error("Failed to create swap chain.");
	}
	else
	{
		Log::Info("Successfully created swap chain.");
	}

	vkGetSwapchainImagesKHR(SwapConfig.Device, SwapChain, &ImageCount, nullptr);
	SwapChainImages.resize(ImageCount);
	vkGetSwapchainImagesKHR(SwapConfig.Device, SwapChain, &ImageCount, SwapChainImages.data());

	SwapChainImageFormat = SurfaceFormat.format;
	SwapChainExtent = Extent;
}

void VulkanSwapChain::CreateImageViews()
{
	SwapChainImageViews.resize(SwapChainImages.size());

	int CreatedViews = 0;
	for (size_t i = 0; i < SwapChainImages.size(); i++)
	{
		VkImageViewCreateInfo CreateInfo{};
		CreateInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
		CreateInfo.image = SwapChainImages[i];

		CreateInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
		CreateInfo.format = SwapChainImageFormat;
		CreateInfo.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
		CreateInfo.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
		CreateInfo.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
		CreateInfo.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;

		CreateInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		CreateInfo.subresourceRange.baseMipLevel = 0;
		CreateInfo.subresourceRange.levelCount = 1;
		CreateInfo.subresourceRange.baseArrayLayer = 0;
		CreateInfo.subresourceRange.layerCount = 1;

		if (vkCreateImageView(SwapConfig.Device, &CreateInfo, nullptr, &SwapChainImageViews[i]) != VK_SUCCESS)
		{
			Log::Error("Failed to create image views.");
		}
		else
		{
			CreatedViews++;
		}
	}

	Log::Info("Successfully created " + std::to_string(CreatedViews) + " image views.");
}