#include "RenderDevice.h"

#include "MacroUtils.h"

#include <SDL3/SDL_log.h>

#include <array>
#include <optional>
#include <span>

#include "Frame.h"
#include "GlobalMemory.h"
#include "MathUtil.h"

RenderDevice::~RenderDevice()
{
	ASSERT(!isInit());
}

// TODO: Failure Handling
RenderDevice* CreateRenderDevice(GlobalMemory* mem, RenderDevice::CreateInfo const& createInfo)
{
	ASSERT(createInfo.window);

	volkInitialize();

	VkInstance instance;
	// Create Instance
	{
		VkApplicationInfo constexpr applicationInfo = {
			.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
			.pNext = nullptr,
			.pApplicationName = "Test",
			.applicationVersion = VK_MAKE_API_VERSION(0, 0, 1, 0),
			.pEngineName = "Blaze",
			.engineVersion = VK_MAKE_API_VERSION(0, 0, 1, 0),
			.apiVersion = VK_API_VERSION_1_3,
		};

		uint32_t instanceExtensionCount;
		char const* const* instanceExtensions = SDL_Vulkan_GetInstanceExtensions(&instanceExtensionCount);

		VkInstanceCreateInfo const instanceCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.pApplicationInfo = &applicationInfo,
			.enabledLayerCount = 0,
			.ppEnabledLayerNames = nullptr,
			.enabledExtensionCount = instanceExtensionCount,
			.ppEnabledExtensionNames = instanceExtensions,
		};

		VK_CHECK(vkCreateInstance(&instanceCreateInfo, nullptr, &instance));
		volkLoadInstance(instance);
	}

	VkSurfaceKHR surface;
	// Create Surface
	ASSERT(SDL_Vulkan_CreateSurface(createInfo.window, instance, nullptr, &surface));

	VkPhysicalDevice physicalDeviceInUse = nullptr;
	VkDevice device = nullptr;
	VmaAllocator gpuAllocator = nullptr;
	std::optional<uint32_t> directQueueFamilyIndex;
	VkQueue directQueue = nullptr;
	// Create Device and Queue
	{
		auto tempAllocStart = mem->getState();

		uint32_t physicalDeviceCount;
		VK_CHECK(vkEnumeratePhysicalDevices(instance, &physicalDeviceCount, nullptr));
		SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "Found %u GPUs", physicalDeviceCount);

		VkPhysicalDevice* physicalDevices = reinterpret_cast<VkPhysicalDevice*>(mem->allocate(sizeof(VkPhysicalDevice) * physicalDeviceCount));
		VK_CHECK(vkEnumeratePhysicalDevices(instance, &physicalDeviceCount, physicalDevices));

		for (VkPhysicalDevice const physicalDevice : std::span{ physicalDevices, physicalDeviceCount })
		{
			auto tempAllocQueueProperties = mem->getState();

			VkPhysicalDeviceProperties properties;
			vkGetPhysicalDeviceProperties(physicalDevice, &properties);

			SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "GPU: %s", properties.deviceName);

			SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "- API Version %d.%d.%d",
				VK_API_VERSION_MAJOR(properties.apiVersion),
				VK_API_VERSION_MINOR(properties.apiVersion),
				VK_API_VERSION_PATCH(properties.apiVersion));

			constexpr static uint32_t API_PATCH_BITS = 0xFFF;
			if ((properties.apiVersion & (~API_PATCH_BITS)) < VK_API_VERSION_1_3)
			{
				continue;
			}

			if (properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_CPU)
			{
				continue;
			}

			uint32_t queueFamilyCount;
			vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, nullptr);
			VkQueueFamilyProperties* queueFamilyProperties = reinterpret_cast<VkQueueFamilyProperties*>(mem->allocate(sizeof(VkQueueFamilyProperties) * queueFamilyCount));
			vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, queueFamilyProperties);

			for (uint32_t queueFamilyIndex = 0; queueFamilyIndex != queueFamilyCount;
				++queueFamilyIndex)
			{
				VkQueueFamilyProperties const& qfp = queueFamilyProperties[queueFamilyIndex];

				bool hasGraphicsSupport = false;
				bool hasComputeSupport = false;
				bool hasTransferSupport = false;
				bool hasPresentSupport = false;

				SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "- Queue [%d]", queueFamilyIndex);
				if (qfp.queueFlags & VK_QUEUE_GRAPHICS_BIT)
				{
					hasGraphicsSupport = true;
					SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "-- Graphic");
				}
				if (qfp.queueFlags & VK_QUEUE_COMPUTE_BIT)
				{
					hasComputeSupport = true;
					SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "-- Compute");
				}
				if (qfp.queueFlags & VK_QUEUE_TRANSFER_BIT)
				{
					hasTransferSupport = true;
					SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "-- Transfer");
				}

				VkBool32 isSurfaceSupported;
				VK_CHECK(vkGetPhysicalDeviceSurfaceSupportKHR(physicalDevice, queueFamilyIndex, surface, &isSurfaceSupported));

				if (isSurfaceSupported)
				{
					hasPresentSupport = true;
					SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "-- Present");
				}

				if (hasGraphicsSupport and hasComputeSupport and hasTransferSupport and hasPresentSupport)
				{
					physicalDeviceInUse = physicalDevice;
					directQueueFamilyIndex = queueFamilyIndex;
					break;
				}

			}

			mem->restoreState(tempAllocQueueProperties);
		}

		ASSERT(physicalDeviceInUse);
		ASSERT(directQueueFamilyIndex.has_value());

		float priority = 1.0f;
		VkDeviceQueueCreateInfo queueCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.queueFamilyIndex = directQueueFamilyIndex.value(),
			.queueCount = 1,
			.pQueuePriorities = &priority,
		};

		VkPhysicalDeviceVulkan13Features constexpr features13 = {
			.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES,
			.pNext = nullptr,
			.synchronization2 = true,
			.dynamicRendering = true,
		};

		VkPhysicalDeviceFeatures features = {
			.depthClamp = true,
			.samplerAnisotropy = true,
		};

		std::array enabledDeviceExtensions = {
			VK_KHR_SWAPCHAIN_EXTENSION_NAME
		};

		VkDeviceCreateInfo const deviceCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
			.pNext = &features13,
			.flags = 0,
			.queueCreateInfoCount = 1,
			.pQueueCreateInfos = &queueCreateInfo,
			.enabledLayerCount = 0,
			.ppEnabledLayerNames = nullptr,
			.enabledExtensionCount = static_cast<uint32_t>(enabledDeviceExtensions.size()),
			.ppEnabledExtensionNames = enabledDeviceExtensions.data(),
			.pEnabledFeatures = &features,
		};

		VK_CHECK(vkCreateDevice(physicalDeviceInUse, &deviceCreateInfo, nullptr, &device));
		volkLoadDevice(device);

		VmaAllocatorCreateInfo allocatorCreateInfo = {
			.flags = 0,
			.physicalDevice = physicalDeviceInUse,
			.device = device,
			.preferredLargeHeapBlockSize = 0,
			.pAllocationCallbacks = nullptr,
			.pDeviceMemoryCallbacks = nullptr,
			.pHeapSizeLimit = nullptr,
			.pVulkanFunctions = nullptr,
			.instance = instance,
			.vulkanApiVersion = VK_API_VERSION_1_3,
			.pTypeExternalMemoryHandleTypes = nullptr,
		};

		VmaVulkanFunctions vkFunctions;
		VK_CHECK(vmaImportVulkanFunctionsFromVolk(&allocatorCreateInfo, &vkFunctions));
		allocatorCreateInfo.pVulkanFunctions = &vkFunctions;

		VK_CHECK(vmaCreateAllocator(&allocatorCreateInfo, &gpuAllocator));

		vkGetDeviceQueue(device, directQueueFamilyIndex.value(), 0, &directQueue);

		mem->restoreState(tempAllocStart);
	}

	// Swapchain creation
	VkExtent2D swapchainExtent = { createInfo.width, createInfo.height };
	VkFormat swapchainFormat = VK_FORMAT_UNDEFINED;
	VkSwapchainKHR swapchain;
	VkImage* swapchainImages;
	VkImageView* swapchainViews;
	uint32_t swapchainImageCount;
	{
		auto tempAllocStart = mem->getState();

		VkSurfaceCapabilitiesKHR capabilities;
		VK_CHECK(vkGetPhysicalDeviceSurfaceCapabilitiesKHR(physicalDeviceInUse, surface, &capabilities));

		// Image Count Calculation
		swapchainImageCount = 3;
		if (capabilities.maxImageCount > 0)
		{
			swapchainImageCount = std::min(swapchainImageCount, capabilities.maxImageCount);
		}
		swapchainImageCount = std::max(swapchainImageCount, capabilities.minImageCount + 1);

		// Image Size calculation
		{
			auto [minWidth, minHeight] = capabilities.minImageExtent;
			auto [maxWidth, maxHeight] = capabilities.maxImageExtent;
			swapchainExtent.width = Clamp(swapchainExtent.width, minWidth, maxWidth);
			swapchainExtent.height = Clamp(swapchainExtent.height, minHeight, maxHeight);
		}

		uint32_t surfaceFormatCount;
		vkGetPhysicalDeviceSurfaceFormatsKHR(physicalDeviceInUse, surface, &surfaceFormatCount, nullptr);
		VkSurfaceFormatKHR* surfaceFormats = reinterpret_cast<VkSurfaceFormatKHR*>(mem->allocate(sizeof(VkSurfaceFormatKHR*) * surfaceFormatCount));
		vkGetPhysicalDeviceSurfaceFormatsKHR(physicalDeviceInUse, surface, &surfaceFormatCount, surfaceFormats);

		VkSurfaceFormatKHR format = {
			.format = VK_FORMAT_UNDEFINED,
			.colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
		};
		for (auto& surfaceFormat : std::span{ surfaceFormats, surfaceFormatCount })
		{
			if (surfaceFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR)
			{
				SDL_LogInfo(SDL_LOG_CATEGORY_GPU, "Color Space SRGB Found %d", surfaceFormat.format);
				if (surfaceFormat.format == VK_FORMAT_R8G8B8A8_SRGB) {
					format = surfaceFormat;
					break;
				}
				if (surfaceFormat.format == VK_FORMAT_B8G8R8A8_SRGB) {
					format = surfaceFormat;
					break;
				}
				if (surfaceFormat.format == VK_FORMAT_R8G8B8A8_UNORM) {
					format = surfaceFormat;
				}
			}
		}
		ASSERT(format.format != VK_FORMAT_UNDEFINED);
		swapchainFormat = format.format;

		uint32_t presentModeCount;
		vkGetPhysicalDeviceSurfacePresentModesKHR(physicalDeviceInUse, surface, &presentModeCount, nullptr);
		VkPresentModeKHR* presentModes = reinterpret_cast<VkPresentModeKHR*>(mem->allocate(sizeof(VkPresentModeKHR*) * presentModeCount));
		vkGetPhysicalDeviceSurfacePresentModesKHR(physicalDeviceInUse, surface, &presentModeCount, presentModes);

		VkPresentModeKHR presentMode = VK_PRESENT_MODE_FIFO_KHR;
		for (VkPresentModeKHR presentModeIter : std::span{ presentModes, presentModeCount })
		{
			if (presentModeIter == VK_PRESENT_MODE_FIFO_RELAXED_KHR)
			{
				presentMode = presentModeIter;
				break;
			}

			if (presentModeIter == VK_PRESENT_MODE_MAILBOX_KHR)
			{
				presentMode = presentModeIter;
			}
		}

		mem->restoreState(tempAllocStart);

		VkSwapchainCreateInfoKHR const swapchainCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
			.pNext = nullptr,
			.flags = 0,
			.surface = surface,
			.minImageCount = swapchainImageCount,
			.imageFormat = format.format,
			.imageColorSpace = format.colorSpace,
			.imageExtent = swapchainExtent,
			.imageArrayLayers = 1,
			.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
			.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE,
			.queueFamilyIndexCount = 0,
			.pQueueFamilyIndices = nullptr,
			.preTransform = VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR,
			.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
			.presentMode = presentMode,
			.clipped = false,
			.oldSwapchain = nullptr,
		};

		VK_CHECK(vkCreateSwapchainKHR(device, &swapchainCreateInfo, nullptr, &swapchain));

		swapchainImageCount = 0;
		vkGetSwapchainImagesKHR(device, swapchain, &swapchainImageCount, nullptr);
		swapchainImages = reinterpret_cast<VkImage*>(mem->allocate(sizeof(VkImage) * swapchainImageCount));
		vkGetSwapchainImagesKHR(device, swapchain, &swapchainImageCount, swapchainImages);

		swapchainViews = reinterpret_cast<VkImageView*>(mem->allocate(sizeof(VkImageView) * swapchainImageCount));
		for (uint32_t i = 0; i != swapchainImageCount; ++i) {
			VkImageViewCreateInfo const viewCreateInfo = {
				.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
				.pNext = nullptr,
				.flags = 0,
				.image = swapchainImages[i],
				.viewType = VK_IMAGE_VIEW_TYPE_2D,
				.format = format.format,
				.components = {
					VK_COMPONENT_SWIZZLE_IDENTITY,
					VK_COMPONENT_SWIZZLE_IDENTITY,
					VK_COMPONENT_SWIZZLE_IDENTITY,
					VK_COMPONENT_SWIZZLE_IDENTITY
				},
				.subresourceRange = {
					.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
					.baseMipLevel = 0,
					.levelCount = 1,
					.baseArrayLayer = 0,
					.layerCount = 1,
				}
			};

			VK_CHECK(vkCreateImageView(device, &viewCreateInfo, nullptr, &swapchainViews[i]));
		}
	}

	// Init frames.

	Frame* frames = reinterpret_cast<Frame*>(mem->allocate(sizeof(Frame) * swapchainImageCount));
	for (uint32_t i = 0; i != swapchainImageCount; ++i)
	{
		new (frames + i) Frame(device, directQueueFamilyIndex.value());
	}

	Byte* allocation = mem->allocate(sizeof(RenderDevice), alignof(RenderDevice));
	return new (allocation) RenderDevice{
		instance, surface, physicalDeviceInUse,
		device, gpuAllocator, directQueue, directQueueFamilyIndex.value(),
		swapchainFormat, swapchainExtent, swapchain, swapchainImages, swapchainViews, frames, swapchainImageCount,
	};
}

inline bool RenderDevice::isInit() const
{
	return instance and device;
}

void RenderDevice::cleanup()
{
	if (not isInit())
		return;

	for (Frame& frame : std::span{ frames, swapchainImageCount })
	{
		frame.cleanup(*this);
	}

	for (auto const& view : std::span{ swapchainViews, swapchainImageCount })
	{
		vkDestroyImageView(device, view, nullptr);
	}

	vkDestroySwapchainKHR(device, Take(swapchain), nullptr);

	vkDestroyDevice(Take(device), nullptr);

	SDL_Vulkan_DestroySurface(instance, Take(surface), nullptr);

	vkDestroyInstance(Take(instance), nullptr);

	volkFinalize();
}

void RenderDevice::waitIdle() const
{
	VK_CHECK(vkDeviceWaitIdle(device));
}

uint32_t RenderDevice::getNumFrames() const
{
	return swapchainImageCount;
}

RenderDevice::RenderDevice(VkInstance const instance, VkSurfaceKHR const surface, VkPhysicalDevice const physicalDeviceInUse,
	VkDevice const device, VmaAllocator gpuAllocator, VkQueue const directQueue, uint32_t const directQueueFamilyIndex,
	VkFormat const swapchainFormat, VkExtent2D const swapchainExtent, VkSwapchainKHR const swapchain, VkImage* swapchainImages,
	VkImageView* swapchainViews, Frame* frames, uint32_t const swapchainImageCount)
	: instance{ instance }
	, surface{ surface }
	, physicalDeviceInUse{ physicalDeviceInUse }
	, device{ device }
	, gpuAllocator{ gpuAllocator }
	, directQueue{ directQueue }
	, directQueueFamilyIndex{ directQueueFamilyIndex }
	, swapchainFormat{ swapchainFormat }
	, swapchainExtent{ swapchainExtent }
	, swapchain{ swapchain }
	, swapchainImages{ swapchainImages }
	, swapchainViews{ swapchainViews }
	, frames{ frames }
	, swapchainImageCount{ swapchainImageCount }
{
}