#include <Orange/Core/Log.h>

#include <Orange/Render/Swapchain.h>
#include <Orange/Render/VulkanHelpers.h>
#include <Orange/Render/Window.h>
#include <vulkan/vulkan_core.h>

namespace orange
{
    static constexpr uint32_t DefaultFramesInFlight = 2;

    static Result<VkSurfaceFormatKHR> ChooseSwapChainFormat(VkPhysicalDevice physicalDevice, VkSurfaceKHR surface)
    {
        SmallVector<VkSurfaceFormatKHR, 32> surfaceFormats;
        if (!VkEnumerate(vkGetPhysicalDeviceSurfaceFormatsKHR, surfaceFormats, physicalDevice, surface))
        {
            log::err("Surface supports no formats");
            return Result<VkSurfaceFormatKHR>::Error(BasicErrorCode::NotFound);
        }

        for (auto& surfaceFormat : surfaceFormats)
        {
            if (surfaceFormat.format == VK_FORMAT_B8G8R8A8_UNORM && surfaceFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR)
                return Result<VkSurfaceFormatKHR>::Success(surfaceFormat);
        }

        return Result<VkSurfaceFormatKHR>::Error(BasicErrorCode::NotFound);
    }

    static VkPresentModeKHR ChoosePresentMode(VkPhysicalDevice physicalDevice, VkSurfaceKHR surface)
    {
        SmallVector<VkPresentModeKHR, 32> presentModes;
        VkEnumerate(vkGetPhysicalDeviceSurfacePresentModesKHR, presentModes, physicalDevice, surface);

        for (auto& presentMode : presentModes) { if (presentMode == VK_PRESENT_MODE_MAILBOX_KHR)   return presentMode; }
        for (auto& presentMode : presentModes) { if (presentMode == VK_PRESENT_MODE_IMMEDIATE_KHR) return presentMode; }

        return VK_PRESENT_MODE_FIFO_KHR;
    }

    static VkExtent2D ChooseSwapExtent(const VkSurfaceCapabilitiesKHR& capabilities)
    {
        if (capabilities.currentExtent.width != ~0u)
            return capabilities.currentExtent;

        uint32_t width = 1280;
        uint32_t height = 720;

        return VkExtent2D
        {
            Clamp(width,  capabilities.minImageExtent.width,  capabilities.maxImageExtent.width),
            Clamp(height, capabilities.minImageExtent.height, capabilities.maxImageExtent.height)
        };
    }

    static uint32_t ChooseImageCount(const VkSurfaceCapabilitiesKHR& capabilities)
    {
        uint32_t imageCount = 3;

        if (capabilities.minImageCount != 0)
            imageCount = Max(imageCount, capabilities.minImageCount);

        if (capabilities.maxImageCount != 0)
            imageCount = Min(imageCount, capabilities.maxImageCount);

        return imageCount;
    }

    Result<Swapchain> Swapchain::Create(RenderContext& context, VkSurfaceKHR surface)
    {
        VkSurfaceCapabilitiesKHR capabilities;
        vkGetPhysicalDeviceSurfaceCapabilitiesKHR(context.PhysicalDevice(), surface, &capabilities);

        Result<VkSurfaceFormatKHR> r_format = ChooseSwapChainFormat(context.PhysicalDevice(), surface);
        if (!r_format)
            return Result<Swapchain>::PrintForwardError(r_format, "Failed to pick swapchain format");

        const VkExtent2D extent = ChooseSwapExtent(capabilities);

        VkSwapchainCreateInfoKHR swapchainInfo =
        {
            .sType            = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
            .surface          = surface,
            .minImageCount    = ChooseImageCount(capabilities),
            .imageFormat      = r_format->format,
            .imageColorSpace  = r_format->colorSpace,
            .imageExtent      = extent,
            .imageArrayLayers = 1,
            .imageUsage       = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
            .imageSharingMode = VK_SHARING_MODE_EXCLUSIVE,
            .preTransform     = capabilities.currentTransform,
            .compositeAlpha   = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
            .presentMode      = ChoosePresentMode(context.PhysicalDevice(), surface),
            .clipped          = VK_TRUE,
            .oldSwapchain     = VK_NULL_HANDLE, // TODO
        };

        VkSwapchainKHR swapchain = VK_NULL_HANDLE;
        if (vkCreateSwapchainKHR(context.Device(), &swapchainInfo, nullptr, &swapchain) != VK_SUCCESS)
            return Result<Swapchain>::PrintError("Failed to create swapchain");

        SmallVector<VkImage, 6> swapchainImages;
        if (!VkEnumerate(vkGetSwapchainImagesKHR, swapchainImages, context.Device(), swapchain))
            return Result<Swapchain>::PrintError("Failed to get swapchain images");

        SmallVector<VkImageView, 6> swapchainImageViews{ swapchainImages.Size() };
        for (size_t i = 0; i < swapchainImages.Size(); i++)
        {
            VkImageViewCreateInfo imageViewInfo =
            {
                .sType    = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
                .image    = swapchainImages[i],
                .viewType = VK_IMAGE_VIEW_TYPE_2D,
                .format   = r_format->format,
                .subresourceRange = FirstMipSubresourceRange,
            };

            if (vkCreateImageView(context.Device(), &imageViewInfo, nullptr, &swapchainImageViews[i]) != VK_SUCCESS)
                return Result<Swapchain>::PrintError("Failed to get swapchain image view");
        }

        SmallVector<VkCommandBuffer, 6> commandBuffers{ DefaultFramesInFlight };
        {
            VkCommandBufferAllocateInfo commandBufferInfo =
            {
                .sType              = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
                .commandPool        = context.CommandPool(),
                .level              = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
                .commandBufferCount = DefaultFramesInFlight,
            };

            if (vkAllocateCommandBuffers(context.Device(), &commandBufferInfo, commandBuffers.Data()) != VK_SUCCESS)
                return Result<Swapchain>::PrintError("Failed to get swapchain image view");
        }

        SmallVector<VkSemaphore, MaxFramesInFlight> imageAvailableSemaphores{ DefaultFramesInFlight };
        SmallVector<VkSemaphore, MaxFramesInFlight> renderFinishedSemaphores{ DefaultFramesInFlight };
        SmallVector<VkFence,     MaxFramesInFlight> inFlightFences          { DefaultFramesInFlight };
        for (uint32_t i = 0; i < DefaultFramesInFlight; i++)
        {
            auto r_fence = context.CreateFence(true);
            if (!r_fence) return Result<Swapchain>::PrintError("Failed to create fence for swapchain");
            inFlightFences[i] = *r_fence;

            auto r_imageSem = context.CreateSemaphore();
            if (!r_imageSem) return Result<Swapchain>::PrintError("Failed to create semaphore for swapchain");
            imageAvailableSemaphores[i] = *r_imageSem;

            auto r_renderSem = context.CreateSemaphore();
            if (!r_renderSem) return Result<Swapchain>::PrintError("Failed to create semaphore for swapchain");
            renderFinishedSemaphores[i] = *r_renderSem;
        }

        // TODO: Handle failure
        uint32_t imageIndex;
        vkAcquireNextImageKHR(context.Device(), swapchain, ~0u, imageAvailableSemaphores[0], VK_NULL_HANDLE, &imageIndex);
    }

    Swapchain::~Swapchain()
    {
        for (const auto& imageView : m_imageViews)
            vkDestroyImageView(m_ctx.Device(), imageView, nullptr);
        vkDestroySwapchainKHR(m_ctx.Device(), m_swapchain, nullptr);
    }
}