/// =============================================
//  Aster: swapchain.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// ==============================================

#include "swapchain.h"

#include "device.h"
#include "physical_device.h"
#include "window.h"

vk::SurfaceCapabilitiesKHR
getSurfaceCapabilities(vk::PhysicalDevice physicalDevice, const vk::SurfaceKHR surface)
{
    vk::SurfaceCapabilitiesKHR surfaceCapabilities;

    vk::Result result = physicalDevice.getSurfaceCapabilitiesKHR(surface, &surfaceCapabilities);
    ERROR_IF(failed(result), "Fetching surface capabilities failed. Cause: {}", result)
    THEN_ABORT(result);

    return surfaceCapabilities;
}

eastl::vector<vk::SurfaceFormatKHR>
getSurfaceFormats(vk::PhysicalDevice physicalDevice, const vk::SurfaceKHR surface)
{
    u32 count = 0;
    vk::Result result = physicalDevice.getSurfaceFormatsKHR(surface, &count, nullptr);
    ERROR_IF(failed(result), "Fetching surface formats failed. Cause: {}", result)
    THEN_ABORT(result);

    eastl::vector<vk::SurfaceFormatKHR> surfaceFormats(count);
    result = physicalDevice.getSurfaceFormatsKHR(surface, &count, surfaceFormats.data());
    ERROR_IF(failed(result), "Fetching surface formats failed. Cause: {}", result)
    THEN_ABORT(result);

    return surfaceFormats;
}

eastl::vector<vk::PresentModeKHR>
getSurfacePresentModes(vk::PhysicalDevice physicalDevice, const vk::SurfaceKHR surface)
{
    u32 count = 0;
    vk::Result result = physicalDevice.getSurfacePresentModesKHR(surface, &count, nullptr);
    ERROR_IF(failed(result), "Fetching surface present modes failed. Cause: {}", result)
    THEN_ABORT(result);

    eastl::vector<vk::PresentModeKHR> surfacePresentModes(count);
    result = physicalDevice.getSurfacePresentModesKHR(surface, &count, surfacePresentModes.data());
    ERROR_IF(failed(result), "Fetching surface present modes failed. Cause: {}", result)
    THEN_ABORT(result);

    return surfacePresentModes;
}

Swapchain::Swapchain(const Window *window, const Device *device, NameString &&name)
    : m_Device(device)
    , m_Name(std::move(name))
{
    this->Create(window);
}

Swapchain::~Swapchain()
{
    if (m_Swapchain)
    {
        m_Device->m_Device.destroy(m_Swapchain, nullptr);
        m_Swapchain = nullptr;
        DEBUG("Swapchain '{}' destroyed.", m_Name);
    }
}

void
Swapchain::Create(const Window *window)
{
    auto surfaceCapabilities = getSurfaceCapabilities(m_Device->m_PhysicalDevice, window->m_Surface);
    auto surfaceFormats = getSurfaceFormats(m_Device->m_PhysicalDevice, window->m_Surface);
    auto presentModes = getSurfacePresentModes(m_Device->m_PhysicalDevice, window->m_Surface);

    vk::Format swapchainFormat = vk::Format::eUndefined;
    vk::ColorSpaceKHR swapchainColorSpace = vk::ColorSpaceKHR::eSrgbNonlinear;
    for (auto [format, colorSpace] : surfaceFormats)
    {
        if (format == vk::Format::eB8G8R8A8Srgb && colorSpace == vk::ColorSpaceKHR::eSrgbNonlinear)
        {
            swapchainFormat = format;
            swapchainColorSpace = colorSpace;
        }
    }
    if (swapchainFormat == vk::Format::eUndefined)
    {
        WARN("Preferred Swapchain format not found. Falling back.");
        auto surfaceFormat = surfaceFormats.front();
        swapchainFormat = surfaceFormat.format;
        swapchainColorSpace = surfaceFormat.colorSpace;
    }

    vk::PresentModeKHR swapchainPresentMode = vk::PresentModeKHR::eFifo;
    for (auto presentMode : presentModes)
    {
        if (presentMode == vk::PresentModeKHR::eMailbox)
        {
            swapchainPresentMode = presentMode;
            break;
        }
    }

    vk::Extent2D swapchainExtent;
    if (surfaceCapabilities.currentExtent.width != MaxValue<u32>)
    {
        swapchainExtent = surfaceCapabilities.currentExtent;
    }
    else
    {
        vk::Extent2D extent = window->GetSize();
        vk::Extent2D minExtent = surfaceCapabilities.minImageExtent;
        vk::Extent2D maxExtent = surfaceCapabilities.maxImageExtent;

        swapchainExtent.width = glm::clamp(extent.width, minExtent.width, maxExtent.width);
        swapchainExtent.height = glm::clamp(extent.height, minExtent.height, maxExtent.height);
    }

    u32 swapchainImageCount = 3;
    if (surfaceCapabilities.maxImageCount > 0)
    {
        swapchainImageCount =
            glm::clamp(swapchainImageCount, surfaceCapabilities.minImageCount, surfaceCapabilities.maxImageCount);
    }

    // TODO: Note that different queues might need the images to be shared.

    const vk::SwapchainCreateInfoKHR swapchainCreateInfo = {
        .surface = window->m_Surface,
        .minImageCount = swapchainImageCount,
        .imageFormat = swapchainFormat,
        .imageColorSpace = swapchainColorSpace,
        .imageExtent = swapchainExtent,
        .imageArrayLayers = 1,
        .imageUsage = vk::ImageUsageFlagBits::eColorAttachment,
        .imageSharingMode = vk::SharingMode::eExclusive,
        .preTransform = surfaceCapabilities.currentTransform,
        .compositeAlpha = vk::CompositeAlphaFlagBitsKHR::eOpaque,
        .presentMode = swapchainPresentMode,
        .clipped = true,
        .oldSwapchain = m_Swapchain,
    };

    vk::SwapchainKHR swapchain;
    vk::Result result = m_Device->m_Device.createSwapchainKHR(&swapchainCreateInfo, nullptr, &swapchain);
    ERROR_IF(failed(result), "Swapchain {} creation failed. Cause {}", m_Name, result)
    THEN_ABORT(result)
    ELSE_DEBUG("Created Swapchain '{}'", m_Name);

    m_Swapchain = swapchain;
}