#include <Orange/Core/Log.h>

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

namespace orange
{
    static constexpr uint32_t DefaultFramesInFlight = 2;

    static Result<VkSurfaceFormatKHR> ChooseSwapChainFormat(VkPhysicalDevice physicalDevice, VkSurfaceKHR surface, bool hdr)
    {
        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 (hdr)
            {
                if (surfaceFormat.format == VK_FORMAT_A2R10G10B10_UNORM_PACK32 && surfaceFormat.colorSpace == VK_COLOR_SPACE_HDR10_ST2084_EXT)
                    return Result<VkSurfaceFormatKHR>::Success(surfaceFormat);

                if (surfaceFormat.format == VK_FORMAT_A2B10G10R10_UNORM_PACK32 && surfaceFormat.colorSpace == VK_COLOR_SPACE_HDR10_ST2084_EXT)
                    return Result<VkSurfaceFormatKHR>::Success(surfaceFormat);
            }
            else
            {
                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 = 800;

        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, bool hdr)
    {
        VkSurfaceCapabilitiesKHR capabilities;
        vkGetPhysicalDeviceSurfaceCapabilitiesKHR(context.PhysicalDevice(), surface, &capabilities);

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

        const VkExtent2D extent = ChooseSwapExtent(capabilities);

        VkFormat viewFormats[2] =
        {
            r_format->format,
            FormatToSrgbFormat(r_format->format),
        };

        VkImageFormatListCreateInfo formatListInfo =
        {
            .sType           = VK_STRUCTURE_TYPE_IMAGE_FORMAT_LIST_CREATE_INFO,
            .viewFormatCount = Size(viewFormats),
            .pViewFormats    = viewFormats,
        };

        VkSwapchainCreateInfoKHR swapchainInfo =
        {
            .sType            = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
            .pNext            = &formatListInfo,
            .flags            = VK_SWAPCHAIN_CREATE_MUTABLE_FORMAT_BIT_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   = FormatToSrgbFormat(r_format->format),
                .subresourceRange = FirstColorMipSubresourceRange,
            };

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

        return Result<Swapchain>::Success(
            context, surface, r_format->format, extent, swapchain,
            swapchainImages, swapchainImageViews, commandBuffers,
            imageAvailableSemaphores, renderFinishedSemaphores, inFlightFences,
            imageIndex);
    }

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

    void Swapchain::Present()
    {
        VkResult res = VK_SUCCESS;

        vkResetFences(m_ctx.Device(), 1, &m_inFlightFences[m_currentFrame]);

        const VkPipelineStageFlags waitStage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
        VkSubmitInfo submitInfo =
        {
            .sType                = VK_STRUCTURE_TYPE_SUBMIT_INFO,
            .waitSemaphoreCount   = 1,
            .pWaitSemaphores      = &m_imageAvailableSemaphores[m_currentFrame],
            .pWaitDstStageMask    = &waitStage,
            .commandBufferCount   = 1,
            .pCommandBuffers      = &m_commandBuffers[m_currentFrame],
            .signalSemaphoreCount = 1,
            .pSignalSemaphores    = &m_renderFinishedSemaphores[m_currentFrame],
        };

        // TODO: Handle failure properly.
        if ((res = vkQueueSubmit(m_ctx.Queue(), 1, &submitInfo, m_inFlightFences[m_currentFrame])) != VK_SUCCESS) {
            log::err("Failed to submit work: %d", res);
            return;
        }

        VkPresentInfoKHR presentInfo =
        {
            .sType              = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
            .waitSemaphoreCount = 1,
            .pWaitSemaphores    = &m_renderFinishedSemaphores[m_currentFrame],
            .swapchainCount     = 1,
            .pSwapchains        = &m_swapchain,
            .pImageIndices      = &m_currentImage,
        };

        res = vkQueuePresentKHR(m_ctx.Queue(), &presentInfo);
        if (res != VK_SUCCESS && res != VK_SUBOPTIMAL_KHR) {
            log::err("Failed to submit present: %d", res);
            return;
        }

        // Prepare for next frame.

        m_currentFrame = (m_currentFrame + 1) % DefaultFramesInFlight;

        vkWaitForFences(m_ctx.Device(), 1, &m_inFlightFences[m_currentFrame], VK_TRUE, UINT64_MAX);
        vkAcquireNextImageKHR(m_ctx.Device(), m_swapchain, ~0u, m_imageAvailableSemaphores[m_currentFrame], VK_NULL_HANDLE, &m_currentImage);

        vkResetCommandBuffer(m_commandBuffers[m_currentFrame], 0u);
    }
}