#include <algorithm>

#include <Vk/Swapchain/Swapchain.hpp>
#include <Utilities/Utilities.hpp>
#include <Log/Log.hpp>
#include <vulkan/vulkan_core.h>

VkSurfaceFormatKHR Vk::Swapchain::chooseSwapSurfaceFormat(const std::vector<VkSurfaceFormatKHR>& availableFormats)
{
    for (const auto& availableFormat : availableFormats)
    {
        if (availableFormat.format == VK_FORMAT_B8G8R8A8_SRGB && availableFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR)
        {
            return availableFormat;
        }
    }

    return availableFormats[0];
}

VkPresentModeKHR Vk::Swapchain::chooseSwapPresentMode(const std::vector<VkPresentModeKHR>& availablePresentModes)
{
    for (const auto& availablePresentMode : availablePresentModes)
    {
        if (availablePresentMode == VK_PRESENT_MODE_MAILBOX_KHR)
        {
            return availablePresentMode;
        }
    }

    return VK_PRESENT_MODE_FIFO_KHR;
}

VkExtent2D Vk::Swapchain::chooseSwapExtent(const VkSurfaceCapabilitiesKHR& capabilities, int width, int height)
{
    if (capabilities.currentExtent.width != std::numeric_limits<uint32_t>::max())
    {
        return capabilities.currentExtent;
    }
    else
    {
        VkExtent2D actualExtent = {
            static_cast<uint32_t>(width),
            static_cast<uint32_t>(height)
        };

        actualExtent.width = std::clamp(actualExtent.width, capabilities.minImageExtent.width, capabilities.maxImageExtent.width);
        actualExtent.height = std::clamp(actualExtent.height, capabilities.minImageExtent.height, capabilities.maxImageExtent.height);

        return actualExtent;
    }
}

Vk::Swapchain::Swapchain(
    std::shared_ptr<Vk::Device> device,
    std::shared_ptr<Vk::Surface> surface,
    int fbWidth, int fbHeight
) :
    m_device(device),
    m_surface(surface)
{
    Vk::Device::SwapchainSupportDetails swapchainSupport = device->getSwapchainSupportDetails();
    Vk::Device::QueueFamilies indices = device->getQueueFamilies();

    VkSurfaceFormatKHR surfaceFormat = chooseSwapSurfaceFormat(swapchainSupport.formats);
    VkPresentModeKHR presentMode = chooseSwapPresentMode(swapchainSupport.presentModes);
    VkExtent2D extent = chooseSwapExtent(swapchainSupport.capabilities, fbWidth, fbHeight);

    uint32_t imageCount = swapchainSupport.capabilities.minImageCount + 1;
    if (swapchainSupport.capabilities.maxImageCount > 0 && imageCount > swapchainSupport.capabilities.maxImageCount)
    {
        imageCount = swapchainSupport.capabilities.maxImageCount;
    }

    VkSwapchainCreateInfoKHR createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
    createInfo.surface = surface->inner();
    createInfo.minImageCount = imageCount;
    createInfo.imageFormat = surfaceFormat.format;
    createInfo.imageColorSpace = surfaceFormat.colorSpace;
    createInfo.imageExtent = extent;
    createInfo.imageArrayLayers = 1;
    createInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT;

    uint32_t queueFamilyIndices[] = {indices.graphicsFamily.value(), indices.presentFamily.value()};
    if (indices.graphicsFamily != indices.presentFamily)
    {
        createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
        createInfo.queueFamilyIndexCount = 2;
        createInfo.pQueueFamilyIndices = queueFamilyIndices;
    }
    else
    {
        createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
        createInfo.queueFamilyIndexCount = 0;
        createInfo.pQueueFamilyIndices = nullptr;
    }

    createInfo.preTransform = swapchainSupport.capabilities.currentTransform;
    createInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
    createInfo.presentMode = presentMode;
    createInfo.clipped = VK_TRUE;
    createInfo.oldSwapchain = VK_NULL_HANDLE;

    if (vkCreateSwapchainKHR(device->inner(), &createInfo, nullptr, &m_swapchain) != VK_SUCCESS)
    {
        crash("Failed to create the swap chain!");
    }

    vkGetSwapchainImagesKHR(device->inner(), m_swapchain, &imageCount, nullptr);
    m_images.resize(imageCount);
    vkGetSwapchainImagesKHR(device->inner(), m_swapchain, &imageCount, m_images.data());

    m_imageFormat = surfaceFormat.format;
    m_extent = extent;
}

Vk::Swapchain::~Swapchain()
{
    vkDestroySwapchainKHR(m_device->inner(), m_swapchain, nullptr);
}

VkSwapchainKHR Vk::Swapchain::inner()
{
    return m_swapchain;
}

VkExtent2D Vk::Swapchain::getExtent()
{
    return m_extent;
}

VkFormat Vk::Swapchain::getImageFormat()
{
    return m_imageFormat;
}

std::vector<VkImage>& Vk::Swapchain::images()
{
    return m_images;
}

std::optional<uint32_t> Vk::Swapchain::acquireNextImage(
    uint64_t timeout,
    Vk::Semaphore *semaphore,
    Vk::Fence *fence
)
{
    uint32_t imageIndex;
    VkSemaphore innerSemaphore = semaphore != nullptr ? semaphore->inner() : VK_NULL_HANDLE;
    VkFence innerFence = fence != nullptr ? fence->inner() : VK_NULL_HANDLE;

    switch (vkAcquireNextImageKHR(m_device->inner(), m_swapchain, timeout, innerSemaphore, innerFence, &imageIndex))
    {
    case VK_SUCCESS:
    case VK_SUBOPTIMAL_KHR:
        return imageIndex;
    default:
        return std::nullopt;
    }
}