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

#include "frame.h"

#include "device.h"
#include "helpers.h"
#include "swapchain.h"

Frame::Frame(const Device *device, const u32 queueFamilyIndex, const u32 frameCount)
    : m_FrameIdx(0)
    , m_ImageIdx(0)
{
    m_Device = device;

    const vk::CommandPoolCreateInfo commandPoolCreateInfo = {
        .flags = vk::CommandPoolCreateFlagBits::eTransient,
        .queueFamilyIndex = queueFamilyIndex,
    };
    AbortIfFailedMV(device->m_Device.createCommandPool(&commandPoolCreateInfo, nullptr, &m_Pool),
                    "Could not command pool for frame {}", frameCount);

    constexpr vk::FenceCreateInfo fenceCreateInfo = {.flags = vk::FenceCreateFlagBits::eSignaled};
    AbortIfFailedMV(device->m_Device.createFence(&fenceCreateInfo, nullptr, &m_FrameAvailableFence),
                    "Could not create a fence for frame {}", frameCount);

    constexpr vk::SemaphoreCreateInfo semaphoreCreateInfo = {};
    AbortIfFailedMV(device->m_Device.createSemaphore(&semaphoreCreateInfo, nullptr, &m_ImageAcquireSem),
                    "Could not create IA semaphore for frame {}.", frameCount);
    AbortIfFailedMV(device->m_Device.createSemaphore(&semaphoreCreateInfo, nullptr, &m_RenderFinishSem),
                    "Could not create RF semaphore for frame {}.", frameCount);

    const vk::CommandBufferAllocateInfo allocateInfo = {
        .commandPool = m_Pool,
        .level = vk::CommandBufferLevel::ePrimary,
        .commandBufferCount = 1,
    };

    AbortIfFailed(m_Device->m_Device.allocateCommandBuffers(&allocateInfo, &m_CommandBuffer));

    DEBUG("Frame {} created successfully.", frameCount);
}

void
Frame::Present(const vk::Queue commandQueue, Swapchain *swapchain, const Window *window)
{

    vk::PresentInfoKHR presentInfo = {
        .waitSemaphoreCount = 1,
        .pWaitSemaphores = &m_RenderFinishSem,
        .swapchainCount = 1,
        .pSwapchains = &swapchain->m_Swapchain,
        .pImageIndices = &m_ImageIdx,
        .pResults = nullptr,
    };
    switch (auto result = commandQueue.presentKHR(&presentInfo))
    {
    case vk::Result::eSuccess:
        break;
    case vk::Result::eErrorOutOfDateKHR:
    case vk::Result::eSuboptimalKHR:
        DEBUG("Recreating Swapchain. Cause: {}", result);
        swapchain->Create(window);
        break; // Present failed. We do nothing. Frame is skipped.
    default:
        AbortIfFailedM(result, "Swapchain Present failed.");
    }
}

Frame::~Frame()
{
    m_Device->m_Device.destroy(m_RenderFinishSem, nullptr);
    m_Device->m_Device.destroy(m_ImageAcquireSem, nullptr);
    m_Device->m_Device.destroy(m_FrameAvailableFence, nullptr);
    m_Device->m_Device.destroy(m_Pool, nullptr);

    DEBUG("Destoryed Frame");
}

// Frame Manager

FrameManager::FrameManager(const Device *device, u32 queueFamilyIndex, u32 framesInFlight)
    : m_Device(device)
    , m_FramesInFlight(framesInFlight)
{
    for (u32 i = 0; i < framesInFlight; ++i)
    {
        m_Frames.emplace_back(device, queueFamilyIndex, i);
    }
}

Frame *
FrameManager::GetNextFrame(Swapchain *swapchain, const Window *window)
{

    Frame *currentFrame = &m_Frames[m_CurrentFrameIdx];
    u32 frameIndex = m_CurrentFrameIdx;

    m_CurrentFrameIdx = (m_CurrentFrameIdx + 1) % m_FramesInFlight;

    AbortIfFailedMV(m_Device->m_Device.waitForFences(1, &currentFrame->m_FrameAvailableFence, true, MaxValue<u64>),
                    "Waiting for fence {} failed.", frameIndex);

    u32 imageIndex = 0;
    bool imageAcquired = false;
    while (!imageAcquired)
    {
        auto result = m_Device->m_Device.acquireNextImageKHR(swapchain->m_Swapchain, MaxValue<u64>,
                                                             currentFrame->m_ImageAcquireSem, nullptr, &imageIndex);
        switch (result)
        {
        case vk::Result::eSuccess:
        case vk::Result::eSuboptimalKHR: // Suboptimal can still render. Better to let this go for semaphores etc.
            imageAcquired = true;
            break; // Image acquired. Break out of loop.
        case vk::Result::eErrorOutOfDateKHR:
            DEBUG("Recreating Swapchain. Cause: {}", result);
            swapchain->Create(window);
            break; // Image acquire has failed. We move to the next frame.
        default:
            AbortIfFailedMV(result, "Waiting for swapchain image {} failed.", frameIndex);
        }
    }

    // Reset fences here. In case swapchain was out of date, we leave the fences signalled.
    AbortIfFailedMV(m_Device->m_Device.resetFences(1, &currentFrame->m_FrameAvailableFence), "Fence {} reset failed.",
                    frameIndex);

    AbortIfFailedMV(m_Device->m_Device.resetCommandPool(currentFrame->m_Pool, {}), "Command pool {} reset failed.",
                    frameIndex);

    currentFrame->m_ImageIdx = imageIndex;

    return currentFrame;
}