#include "VulkanDeviceManager.h"

#include <vector>
#include <map>
#include <cstdint> // Necessary for uint32_t
#include <set>
#include <vulkan/vulkan_core.h>

#include "VulkanContext.hpp"
#include "utilities/Logger.h"

const std::vector<const char*> DeviceExtensions = {
	VK_KHR_SWAPCHAIN_EXTENSION_NAME
};

VulkanDeviceManager::VulkanDeviceManager(
	VkInstance	 Instance,
	VkSurfaceKHR Surface,
	GLFWwindow*	 Window,
	bool		 bEnableValidationLayers)
	: Instance(Instance), Surface(Surface), Window(Window), bEnableValidationLayers(bEnableValidationLayers)
{
}

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

void VulkanDeviceManager::Cleanup()
{
	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();
	}

	VkPhysicalDeviceFeatures SupportedFeatures;
	vkGetPhysicalDeviceFeatures(Device, &SupportedFeatures);

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

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, Surface, &PresentSupport);

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

		i++;
	}

	return Indices;
}

void VulkanDeviceManager::CreateLogicalDevice()
{
	PhysicalQueueFamilies = FindQueueFamilies(PhysicalDevice);

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

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

	VkPhysicalDeviceFeatures DeviceFeatures{};
	DeviceFeatures.samplerAnisotropy = VK_TRUE;

	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!");
	}
	else
	{
		Log::Info("Created logical device.");
	}

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

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

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

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

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

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

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

	return Details;
}
