#include "VulkanDeviceManager.h"

#include <vector>
#include <map>
#include <cstdint>	 // Necessary for uint32_t
#include <limits>	 // Necessary for std::numeric_limits
#include <algorithm> // Necessary for std::clamp

#include "Logger.h"
#include "GlfwWindowManager.h"
#include <set>

std::vector<VkImage> VulkanDeviceManager::SwapChainImages = {};

VulkanDeviceManager::VulkanDeviceManager()
{
}

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

VulkanDeviceManager::VulkanDeviceManager(VulkanDeviceManager&& Other) noexcept
	: PhysicalDevice(std::exchange(Other.PhysicalDevice, VK_NULL_HANDLE))
	, Instance(std::exchange(Other.Instance, VK_NULL_HANDLE))
	, Device(std::exchange(Other.Device, VK_NULL_HANDLE))
	, GraphicsQueue(std::exchange(Other.GraphicsQueue, VK_NULL_HANDLE))
	, bEnableValidationLayers(std::exchange(Other.bEnableValidationLayers, false))
	, ValidationLayers(std::move(Other.ValidationLayers))
{
}

VulkanDeviceManager& VulkanDeviceManager::operator=(VulkanDeviceManager&& Other) noexcept
{
	if (this != &Other)
	{
		Cleanup(); // Clean up current resources

		// Transfer resources from Other
		PhysicalDevice = std::exchange(Other.PhysicalDevice, VK_NULL_HANDLE);
		Instance = std::exchange(Other.Instance, VK_NULL_HANDLE);
		Device = std::exchange(Other.Device, VK_NULL_HANDLE);
		GraphicsQueue = std::exchange(Other.GraphicsQueue, VK_NULL_HANDLE);
		bEnableValidationLayers = std::exchange(Other.bEnableValidationLayers, false);
		ValidationLayers = std::move(Other.ValidationLayers);
	}
	return *this;
}

void VulkanDeviceManager::Initialize(
	VkInstance						Instance,
	bool							bEnableValidationLayers,
	const std::vector<const char*>& ValidationLayers)
{
	if (IsInitialized())
	{
		Log::Warning("Already Initialized.");
		return;
	}

	this->Instance = Instance;
	this->bEnableValidationLayers = bEnableValidationLayers;
	this->ValidationLayers = &ValidationLayers;

	PickPhysicalDevice();
	CreateLogicalDevice();
	CreateSwapChain();
	CreateImageViews();
}

void VulkanDeviceManager::Cleanup()
{
	if (!IsInitialized())
	{
		Log::Warning("Not Initialized.");
		return;
	}
	for (auto ImageView : SwapChainImageViews)
	{
		vkDestroyImageView(Device, ImageView, nullptr);
	}
	vkDestroySwapchainKHR(Device, SwapChain, nullptr);

	vkDestroySurfaceKHR(Instance, GlfwWindowManager::Surface, nullptr);
	vkDestroyDevice(Device, nullptr);
}

void VulkanDeviceManager::PickPhysicalDevice()
{
	uint32_t DeviceCount = 0;
	vkEnumeratePhysicalDevices(Instance, &DeviceCount, nullptr);

	if (DeviceCount == 0)
	{
		Log::Error("Failed to find GPU with Vulkan Support.");
	}

	std::vector<VkPhysicalDevice> Devices(DeviceCount);
	vkEnumeratePhysicalDevices(Instance, &DeviceCount, Devices.data());

	std::multimap<int, VkPhysicalDevice> Candidates;

	for (const auto& Device : Devices)
	{
		if (IsDeviceSuitable(Device))
		{
			int Score = RateDeviceSuitability(Device);
			Candidates.insert(std::make_pair(Score, Device));
		}
	}

	if (Candidates.rbegin()->first > 0)
	{
		PhysicalDevice = Candidates.rbegin()->second;
		Log::Info("Suitable GPU found.");
	}
	else
	{
		Log::Error("Failed to find a suitable GPU.");
	}
}

bool VulkanDeviceManager::IsDeviceSuitable(VkPhysicalDevice Device)
{
	QueueFamilyIndices Indices = FindQueueFamilies(Device);

	bool bExtensionsSupported = CheckDeviceExtensionSupport(Device);

	bool bSwapChainAdequate = false;
	if (bExtensionsSupported)
	{
		SwapChainSupportDetails SwapChainSupport = QuerySwapChainSupport(Device);
		bSwapChainAdequate = !SwapChainSupport.Formats.empty() && !SwapChainSupport.PresentModes.empty();
	}

	return Indices.IsComplete() && bExtensionsSupported && bSwapChainAdequate;
}

int VulkanDeviceManager::RateDeviceSuitability(VkPhysicalDevice Device)
{
	VkPhysicalDeviceProperties DeviceProperties;
	vkGetPhysicalDeviceProperties(Device, &DeviceProperties);

	VkPhysicalDeviceFeatures DeviceFeatures;
	vkGetPhysicalDeviceFeatures(Device, &DeviceFeatures);

	int Score = 0;

	if (DeviceProperties.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU)
	{
		Score += 100;
	}

	Score += DeviceProperties.limits.maxImageDimension2D;

	if (!DeviceFeatures.geometryShader)
	{
		return 0;
	}

	return Score;
}

bool VulkanDeviceManager::CheckDeviceExtensionSupport(VkPhysicalDevice Device)
{
	uint32_t ExtensionCount;
	vkEnumerateDeviceExtensionProperties(Device, nullptr, &ExtensionCount, nullptr);

	std::vector<VkExtensionProperties> AvailableExtensions(ExtensionCount);
	vkEnumerateDeviceExtensionProperties(Device, nullptr, &ExtensionCount, AvailableExtensions.data());

	std::set<std::string> RequiredExtensions(DeviceExtensions.begin(), DeviceExtensions.end());

	for (const auto& Extension : AvailableExtensions)
	{
		RequiredExtensions.erase(Extension.extensionName);
	}

	return RequiredExtensions.empty();
}

QueueFamilyIndices VulkanDeviceManager::FindQueueFamilies(VkPhysicalDevice Device)
{
	QueueFamilyIndices Indices;

	uint32_t QueueFamilyCount = 0;
	vkGetPhysicalDeviceQueueFamilyProperties(Device, &QueueFamilyCount, nullptr);

	std::vector<VkQueueFamilyProperties> QueueFamilies(QueueFamilyCount);
	vkGetPhysicalDeviceQueueFamilyProperties(Device, &QueueFamilyCount, QueueFamilies.data());

	int i = 0;
	for (const auto& QueueFamily : QueueFamilies)
	{
		if (QueueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT)
		{
			Indices.GraphicsFamily = i;
		}

		VkBool32 PresentSupport = false;
		vkGetPhysicalDeviceSurfaceSupportKHR(Device, i, GlfwWindowManager::Surface, &PresentSupport);

		if (PresentSupport)
		{
			Indices.PresentFamily = i;
		}
		if (Indices.IsComplete())
		{
			break;
		}

		i++;
	}

	return Indices;
}

void VulkanDeviceManager::CreateLogicalDevice()
{
	QueueFamilyIndices Indices = FindQueueFamilies(PhysicalDevice);

	std::vector<VkDeviceQueueCreateInfo> QueueCreateInfos;
	std::set<uint32_t>					 UniqueQueueFamilies = { Indices.GraphicsFamily.value(), Indices.PresentFamily.value() };

	float QueuePriority = 1.0f;
	for (uint32_t QueueFamily : UniqueQueueFamilies)
	{
		VkDeviceQueueCreateInfo QueueCreateInfo{};
		QueueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
		QueueCreateInfo.queueFamilyIndex = Indices.GraphicsFamily.value();
		QueueCreateInfo.queueCount = 1;
		QueueCreateInfo.pQueuePriorities = &QueuePriority;
		QueueCreateInfos.push_back(QueueCreateInfo);
	}

	VkPhysicalDeviceFeatures DeviceFeatures{};

	VkDeviceCreateInfo CreateInfo{};
	CreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
	CreateInfo.queueCreateInfoCount = static_cast<uint32_t>(QueueCreateInfos.size());
	CreateInfo.pQueueCreateInfos = QueueCreateInfos.data();
	CreateInfo.pEnabledFeatures = &DeviceFeatures;

	CreateInfo.enabledExtensionCount = static_cast<uint32_t>(DeviceExtensions.size());
	CreateInfo.ppEnabledExtensionNames = DeviceExtensions.data();

	if (bEnableValidationLayers)
	{
		CreateInfo.enabledLayerCount = static_cast<uint32_t>(ValidationLayers->size());
		CreateInfo.ppEnabledLayerNames = ValidationLayers->data();
	}
	else
	{
		CreateInfo.enabledLayerCount = 0;
	}

	if (vkCreateDevice(PhysicalDevice, &CreateInfo, nullptr, &Device) != VK_SUCCESS)
	{
		Log::Error("Failed to create logical device!");
	}

	vkGetDeviceQueue(Device, Indices.GraphicsFamily.value(), 0, &GraphicsQueue);
	vkGetDeviceQueue(Device, Indices.PresentFamily.value(), 0, &PresentQueue);
}

SwapChainSupportDetails VulkanDeviceManager::QuerySwapChainSupport(VkPhysicalDevice Device)
{
	SwapChainSupportDetails Details;

	vkGetPhysicalDeviceSurfaceCapabilitiesKHR(Device, GlfwWindowManager::Surface, &Details.Capabilities);

	uint32_t FormatCount;
	vkGetPhysicalDeviceSurfaceFormatsKHR(Device, GlfwWindowManager::Surface, &FormatCount, nullptr);

	if (FormatCount != 0)
	{
		Details.Formats.resize(FormatCount);
		vkGetPhysicalDeviceSurfaceFormatsKHR(Device, GlfwWindowManager::Surface, &FormatCount, Details.Formats.data());
	}

	uint32_t PresentModeCount;
	vkGetPhysicalDeviceSurfacePresentModesKHR(Device, GlfwWindowManager::Surface, &PresentModeCount, nullptr);

	if (PresentModeCount != 0)
	{
		Details.PresentModes.resize(PresentModeCount);
		vkGetPhysicalDeviceSurfacePresentModesKHR(Device, GlfwWindowManager::Surface, &PresentModeCount, Details.PresentModes.data());
	}

	return Details;
}

VkSurfaceFormatKHR VulkanDeviceManager::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 VulkanDeviceManager::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 VulkanDeviceManager::ChooseSwapExtent(const VkSurfaceCapabilitiesKHR& Capabilities)
{
	if (Capabilities.currentExtent.width != (std::numeric_limits<uint32_t>::max)())
	{
		return Capabilities.currentExtent;
	}
	else
	{
		int Width, Height;
		glfwGetFramebufferSize(GlfwWindowManager::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 VulkanDeviceManager::CreateSwapChain()
{
	SwapChainSupportDetails SwapChainSupport = QuerySwapChainSupport(PhysicalDevice);

	VkSurfaceFormatKHR SurfaceFormat = ChooseSwapSurfaceFormat(SwapChainSupport.Formats);
	VkPresentModeKHR   PresentMode = ChooseSwapPresentMode(SwapChainSupport.PresentModes);
	VkExtent2D		   Extent = ChooseSwapExtent(SwapChainSupport.Capabilities);

	uint32_t ImageCount = SwapChainSupport.Capabilities.minImageCount + 1;

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

	VkSwapchainCreateInfoKHR CreateInfo{};
	CreateInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
	CreateInfo.surface = GlfwWindowManager::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

	QueueFamilyIndices Indices = FindQueueFamilies(PhysicalDevice);
	uint32_t		   QueueFamilyIndices[] = { Indices.GraphicsFamily.value(),
				  Indices.PresentFamily.value() };

	if (Indices.GraphicsFamily != Indices.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 = SwapChainSupport.Capabilities.currentTransform;
	CreateInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
	CreateInfo.presentMode = PresentMode;
	CreateInfo.clipped = VK_TRUE;
	CreateInfo.oldSwapchain = VK_NULL_HANDLE;

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

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

	SwapChainImageFormat = SurfaceFormat.format;
	SwapChainExtent = Extent;
}

void VulkanDeviceManager::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(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.");
}