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

#include "swapchain.h"

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

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

Swapchain::~Swapchain()
{
    this->Cleanup();
}

Swapchain::Swapchain(Swapchain &&other) noexcept
    : m_Device(other.m_Device)
    , m_Swapchain(Take(other.m_Swapchain))
    , m_Name(std::move(other.m_Name))
    , m_Extent(other.m_Extent)
    , m_Format(other.m_Format)
    , m_Images(std::move(other.m_Images))
    , m_ImageViews(std::move(other.m_ImageViews))
{
}

Swapchain &
Swapchain::operator=(Swapchain &&other) noexcept
{
    if (this == &other)
        return *this;
    m_Device = other.m_Device;
    m_Swapchain = Take(other.m_Swapchain);
    m_Name = std::move(other.m_Name);
    m_Extent = other.m_Extent;
    m_Format = other.m_Format;
    m_Images = std::move(other.m_Images);
    m_ImageViews = std::move(other.m_ImageViews);
    return *this;
}

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

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

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

    if (surfaceCapabilities.currentExtent.width != MaxValue<u32>)
    {
        m_Extent = surfaceCapabilities.currentExtent;
    }
    else
    {
        auto [width, height] = window->GetSize();
        auto [minWidth, minHeight] = surfaceCapabilities.minImageExtent;
        auto [maxWidth, maxHeight] = surfaceCapabilities.maxImageExtent;

        m_Extent.width = glm::clamp(width, minWidth, maxWidth);
        m_Extent.height = glm::clamp(height, minHeight, maxHeight);
    }

    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 = m_Format,
        .imageColorSpace = swapchainColorSpace,
        .imageExtent = m_Extent,
        .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::Device device = m_Device->m_Device;

    vk::SwapchainKHR swapchain;
    vk::Result result = 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);

    // Irrelevant on the first run. Required for re-creation.
    Cleanup();

    m_Swapchain = swapchain;

    m_Device->SetName(m_Swapchain, m_Name.data());

    result = device.getSwapchainImagesKHR(m_Swapchain, &swapchainImageCount, nullptr);
    ERROR_IF(Failed(result), "Failed getting swapchain {}'s images. Cause {}", m_Name, result)
    THEN_ABORT(result);

    // Managed by the Swapchain.
    m_Images.resize(swapchainImageCount);
    result = device.getSwapchainImagesKHR(m_Swapchain, &swapchainImageCount, m_Images.data());
    ERROR_IF(Failed(result), "Failed getting swapchain {}'s images. Cause {}", m_Name, result)
    THEN_ABORT(result);

    vk::ImageViewCreateInfo viewCreateInfo = {
        .viewType = vk::ImageViewType::e2D,
        .format = m_Format,
        .components = {.r = vk::ComponentSwizzle::eIdentity,
                       .g = vk::ComponentSwizzle::eIdentity,
                       .b = vk::ComponentSwizzle::eIdentity,
                       .a = vk::ComponentSwizzle::eIdentity},
        .subresourceRange = {.aspectMask = vk::ImageAspectFlagBits::eColor,
                             .baseMipLevel = 0,
                             .levelCount = 1,
                             .baseArrayLayer = 0,
                             .layerCount = 1},
    };

    u32 index = 0;
    for (auto image : m_Images)
    {
        viewCreateInfo.image = image;

        vk::ImageView imageView;
        result = device.createImageView(&viewCreateInfo, nullptr, &imageView);
        ERROR_IF(Failed(result), "Failed creating swapchain {}'s image view [{}]. Cause {}", m_Name, index, result)
        THEN_ABORT(result);

        m_ImageViews.push_back(imageView);

        ++index;
    }

    DEBUG("Swapchain {} Image Views created.", m_Name);
}

void
Swapchain::Cleanup()
{
    if (!m_ImageViews.empty()) // Don't want the condition in the logs.
        DEBUG("Swapchain {} Image Views destroyed.", m_Name);
    for (const auto imageView : m_ImageViews)
    {
        m_Device->m_Device.destroy(imageView, nullptr);
    }
    m_ImageViews.clear();
    if (m_Swapchain)
    {
        m_Device->m_Device.destroy(m_Swapchain, nullptr);
        m_Swapchain = nullptr;
        DEBUG("Swapchain '{}' destroyed.", m_Name);
    }
}