#include "RenderDevice.h"

#include "MacroUtils.h"

#include <SDL3/SDL_log.h>

#include <array>

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

RenderDevice::RenderDevice(CreateInfo const& createInfo)
{
	ASSERT(createInfo.window);

	volkInitialize();

	// 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);
	}

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

	// Create Device and Queue
	{
		uint32_t physicalDeviceCount;
		VK_CHECK(vkEnumeratePhysicalDevices(instance, &physicalDeviceCount, nullptr));
		SDL_Log("Found %u GPUs", physicalDeviceCount);

		std::vector<VkPhysicalDevice> physicalDevices(physicalDeviceCount);
		VK_CHECK(vkEnumeratePhysicalDevices(instance, &physicalDeviceCount, physicalDevices.data()));


		std::vector<VkQueueFamilyProperties> queueFamilyProperties;
		for (VkPhysicalDevice const physicalDevice : physicalDevices)
		{
			VkPhysicalDeviceProperties properties;
			vkGetPhysicalDeviceProperties(physicalDevice, &properties);

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

			SDL_Log("- 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);
			queueFamilyProperties.resize(queueFamilyCount);
			vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, queueFamilyProperties.data());

			for (uint32_t queueFamilyIndex = 0; VkQueueFamilyProperties const qfp : queueFamilyProperties)
			{
				bool hasGraphicsSupport = false;
				bool hasComputeSupport = false;
				bool hasTransferSupport = false;
				bool hasPresentSupport = false;

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

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

				if (isSurfaceSupported)
				{
					hasPresentSupport = true;
					SDL_Log("-- Present");
				}

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

				++queueFamilyIndex;
			}
		}

		ASSERT(physicalDeviceInUse);

		float priority = 1.0f;
		VkDeviceQueueCreateInfo queueCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
			.pNext = nullptr,
			.flags = 0,
			.queueFamilyIndex = directQueueFamilyIndex,
			.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));
		vkGetDeviceQueue(device, directQueueFamilyIndex, 0, &directQueue);
	}

	// Swapchain creation
	swapchainExtent = { createInfo.width, createInfo.height };
	{
		VkSurfaceCapabilitiesKHR capabilities;
		VK_CHECK(vkGetPhysicalDeviceSurfaceCapabilitiesKHR(physicalDeviceInUse, surface, &capabilities));

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

		// 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);
		std::vector<VkSurfaceFormatKHR> surfaceFormats(surfaceFormatCount);
		vkGetPhysicalDeviceSurfaceFormatsKHR(physicalDeviceInUse, surface, &surfaceFormatCount, surfaceFormats.data());

		VkSurfaceFormatKHR format = {
			.format = VK_FORMAT_UNDEFINED,
			.colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
		};
		for (auto& surfaceFormat : surfaceFormats)
		{
			if (surfaceFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR)
			{
				SDL_Log("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);
		std::vector<VkPresentModeKHR> presentModes(presentModeCount);
		vkGetPhysicalDeviceSurfacePresentModesKHR(physicalDeviceInUse, surface, &presentModeCount, presentModes.data());

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

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

		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);
		SDL_Log("%llu", swapchainImages.size());
		swapchainImages.resize(swapchainImageCount);
		vkGetSwapchainImagesKHR(device, swapchain, &swapchainImageCount, swapchainImages.data());

		swapchainViews.reserve(swapchainImageCount);
		for (auto& image : swapchainImages) {

			VkImageViewCreateInfo const viewCreateInfo = {
				.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
				.pNext = nullptr,
				.flags = 0,
				.image = image,
				.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,
				}
			};

			VkImageView view;
			VK_CHECK(vkCreateImageView(device, &viewCreateInfo, nullptr, &view));
			swapchainViews.push_back(view);
		}
	}

	// Init frames.
	{
		auto count = static_cast<uint32_t>(swapchainImages.size());
		frames.reserve(count);
		for (uint32_t i = 0; i != count; ++i)
		{
			frames.emplace_back(*this);
		}
	}
}

RenderDevice::~RenderDevice()
{
	cleanup();
}

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

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

	for (Frame& frame : frames)
	{
		frame.cleanup(*this);
	}
	frames.clear();

	for (auto const& view : swapchainViews)
	{
		vkDestroyImageView(device, view, nullptr);
	}
	swapchainViews.clear();
	swapchainImages.clear();

	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 static_cast<uint32_t>(frames.size());
}