#include <Orange/Core/Log.h>
#include <Orange/Core/Vector.h>

#include <Orange/Render/VulkanHelpers.h>
#include <Orange/Render/RenderContext.h>
#include <Orange/Render/Window.h>

namespace orange
{
    static VulkanResult<VkInstance> CreateVkInstance(const char* appName)
    {
        SmallVector<const char *, 32> instanceExtensions;
        {
            auto r_tempWindow = Window::Create();
            if (!r_tempWindow)
                return VulkanResult<VkInstance>::PrintError(
                    "Failed to create temporary window for grabbing instance extensions: %s", Window::GetError());

            if (!VkEnumerate(Window::GetInstanceExtensions, instanceExtensions, *r_tempWindow))
                return VulkanResult<VkInstance>::PrintError(
                    "Failed to get required instance extensions: %s", Window::GetError());
        }

        VkApplicationInfo appInfo =
        {
            .sType              = VK_STRUCTURE_TYPE_APPLICATION_INFO,
            .pApplicationName   = appName,
            .applicationVersion = VK_MAKE_VERSION(1, 0, 0),
            .pEngineName        = "Orange",
            .engineVersion      = VK_MAKE_VERSION(1, 0, 0),
            .apiVersion         = VK_API_VERSION_1_3,
        };

        VkInstanceCreateInfo instanceInfo =
        {
            .sType                   = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
            .pApplicationInfo        = &appInfo,
            .enabledExtensionCount   = uint32_t(instanceExtensions.Size()),
            .ppEnabledExtensionNames = instanceExtensions.Data(),
        };

        VkResult result = VK_SUCCESS;
        VkInstance instance = VK_NULL_HANDLE;
        if ((result = vkCreateInstance(&instanceInfo, nullptr, &instance)) != VK_SUCCESS)
            return VulkanResult<VkInstance>::PrintError(result, "Failed to create Vulkan instance.");

        return VulkanResult<VkInstance>::Success(instance);
    }

    static Result<uint32_t> PickQueueFamilyIndex(VkPhysicalDevice physicalDevice)
    {
        SmallVector<VkQueueFamilyProperties, 32> queueFamilyProperties;
        VkEnumerate(vkGetPhysicalDeviceQueueFamilyProperties, queueFamilyProperties, physicalDevice);

        for (uint32_t i = 0; i < queueFamilyProperties.Size(); i++)
        {
            const auto& family = queueFamilyProperties[i];
            if (family.queueFlags & VK_QUEUE_GRAPHICS_BIT)
                return Result<uint32_t>::Success(i);
        }
        return Result<uint32_t>::Error(BasicErrorCode::NotFound);
    }

    static int ScorePhysicalDevice(VkPhysicalDevice physicalDevice)
    {
        VkPhysicalDeviceProperties2 props = { VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2 };
        vkGetPhysicalDeviceProperties2(physicalDevice, &props);

        int score = 0;

        if (props.properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU ||
            props.properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU)
            score += 10;

        if (props.properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU)
            score += 5;

        return score;
    }

    static Result<VkPhysicalDevice> PickPhysicalDevice(VkInstance instance)
    {
        SmallVector<VkPhysicalDevice, 32> physicalDevices;
        VkEnumerate(vkEnumeratePhysicalDevices, physicalDevices, instance);

        VkPhysicalDevice bestPhysicalDevice = VK_NULL_HANDLE;
        int bestScore = -1;
        for (auto physicalDevice : physicalDevices)
        {
            int score = ScorePhysicalDevice(physicalDevice);
            if (score > bestScore && PickQueueFamilyIndex(physicalDevice))
            {
                bestPhysicalDevice = physicalDevice;
                bestScore = score;
            }
        }

        if (!bestPhysicalDevice)
            return Result<VkPhysicalDevice>::PrintError("Failed to find a suitable physical device");

        return Result<VkPhysicalDevice>::Success(bestPhysicalDevice);
    }

    static VulkanResult<VkDevice> CreateVkDevice(VkPhysicalDevice physicalDevice, uint32_t graphicsQueueFamilyIndex)
    {
        constexpr float queuePriority = 1.0f;
        VkDeviceQueueCreateInfo queueInfo =
        {
            .sType              = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
            .queueFamilyIndex   = graphicsQueueFamilyIndex,
            .queueCount         = 1,
            .pQueuePriorities   = &queuePriority,
        };

        const char* deviceExtensions[] =
        {
            VK_KHR_SWAPCHAIN_EXTENSION_NAME,
            VK_KHR_SWAPCHAIN_MUTABLE_FORMAT_EXTENSION_NAME,
        };

        VkPhysicalDeviceVulkan12Features vulkan12features =
        {
            .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES,
            .drawIndirectCount = VK_TRUE,
        };

        VkPhysicalDeviceDynamicRenderingFeatures dynamicRenderingFeatures =
        {
            .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DYNAMIC_RENDERING_FEATURES,
            .pNext = &vulkan12features,
            .dynamicRendering = VK_TRUE,
        };

        VkPhysicalDeviceDescriptorIndexingFeatures descriptorIndexingFeatures =
        {
            .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_FEATURES,
            .pNext = &dynamicRenderingFeatures,
            .shaderInputAttachmentArrayDynamicIndexing          = VK_FALSE,
            .shaderUniformTexelBufferArrayDynamicIndexing       = VK_TRUE,
            .shaderStorageTexelBufferArrayDynamicIndexing       = VK_TRUE,
            .shaderUniformBufferArrayNonUniformIndexing         = VK_TRUE,
            .shaderSampledImageArrayNonUniformIndexing          = VK_TRUE,
            .shaderStorageBufferArrayNonUniformIndexing         = VK_TRUE,
            .shaderStorageImageArrayNonUniformIndexing          = VK_TRUE,
            .shaderInputAttachmentArrayNonUniformIndexing       = VK_FALSE,
            .shaderUniformTexelBufferArrayNonUniformIndexing    = VK_TRUE,
            .shaderStorageTexelBufferArrayNonUniformIndexing    = VK_TRUE,
            .descriptorBindingUniformBufferUpdateAfterBind      = VK_TRUE,
            .descriptorBindingSampledImageUpdateAfterBind       = VK_TRUE,
            .descriptorBindingStorageImageUpdateAfterBind       = VK_TRUE,
            .descriptorBindingStorageBufferUpdateAfterBind      = VK_TRUE,
            .descriptorBindingUniformTexelBufferUpdateAfterBind = VK_TRUE,
            .descriptorBindingStorageTexelBufferUpdateAfterBind = VK_TRUE,
            .descriptorBindingUpdateUnusedWhilePending          = VK_TRUE,
            .descriptorBindingPartiallyBound                    = VK_TRUE,
            .descriptorBindingVariableDescriptorCount           = VK_TRUE,
            .runtimeDescriptorArray                             = VK_TRUE,
        };

        VkPhysicalDeviceFeatures features
        {
            .samplerAnisotropy = VK_TRUE,
        };

        VkDeviceCreateInfo deviceInfo =
        {
            .sType                      = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
            .pNext                      = &descriptorIndexingFeatures,
            .queueCreateInfoCount       = 1,
            .pQueueCreateInfos          = &queueInfo,
            .enabledExtensionCount      = Size(deviceExtensions),
            .ppEnabledExtensionNames    = deviceExtensions,
            .pEnabledFeatures           = &features,
        };

        VkDevice device = VK_NULL_HANDLE;
        VkResult result;
        if ((result = vkCreateDevice(physicalDevice, &deviceInfo, nullptr, &device)) != VK_SUCCESS)
            return VulkanResult<VkDevice>::PrintError(result, "Failed to create Vulkan logical device");

        return VulkanResult<VkDevice>::Success(device);
    }

    static VulkanResult<VkCommandPool> CreateVkCommandPool(VkDevice device, uint32_t graphicsQueueFamilyIndex)
    {
        VkCommandPoolCreateInfo poolInfo =
        {
            .sType            = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
            .flags            = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
            .queueFamilyIndex = graphicsQueueFamilyIndex,
        };

        VkResult res = VK_SUCCESS;
        VkCommandPool commandPool = VK_NULL_HANDLE;
        if ((res = vkCreateCommandPool(device, &poolInfo, nullptr, &commandPool)) != VK_SUCCESS)
            return VulkanResult<VkCommandPool>::PrintError(res, "Failed to create command pool");

        return VulkanResult<VkCommandPool>::Success(commandPool);
    }

    VulkanResult<RenderContext> RenderContext::Create(const char *appName)
    {
        auto r_instance = CreateVkInstance(appName);
        if (!r_instance)
            return VulkanResult<RenderContext>::ForwardError(r_instance);

        auto r_physicalDevice = PickPhysicalDevice(*r_instance);
        if (!r_physicalDevice)
            return VulkanResult<RenderContext>::Error();

        uint32_t graphicsQueueFamilyIndex = *PickQueueFamilyIndex(*r_physicalDevice);

        auto r_device = CreateVkDevice(*r_physicalDevice, graphicsQueueFamilyIndex);
        if (!r_device)
            return VulkanResult<RenderContext>::ForwardError(r_device);

        auto r_commandPool = CreateVkCommandPool(*r_device, graphicsQueueFamilyIndex);
        if (!r_commandPool)
            return VulkanResult<RenderContext>::ForwardError(r_commandPool);

        VkQueue queue = VK_NULL_HANDLE;
        vkGetDeviceQueue(*r_device, graphicsQueueFamilyIndex, 0, &queue);

        return VulkanResult<RenderContext>::Success(*r_instance, *r_physicalDevice, *r_device, queue, *r_commandPool);
    }

    RenderContext::~RenderContext()
    {
        vkDestroyCommandPool(m_device, m_commandPool, nullptr);
        vkDestroyDevice(m_device, nullptr);
        vkDestroyInstance(m_instance, nullptr);
    }

}