#include "VKCommandManager.hpp"

#include "MayaFlux/Journal/Archivist.hpp"

namespace MayaFlux::Core {

VKCommandManager::~VKCommandManager()
{
    cleanup();
}

bool VKCommandManager::initialize(vk::Device device, uint32_t graphics_queue_family, uint32_t compute_queue_family)
{
    m_device = device;
    m_graphics_queue_family = graphics_queue_family;
    m_compute_queue_family = compute_queue_family;

    vk::CommandPoolCreateInfo pool_info {};
    pool_info.queueFamilyIndex = graphics_queue_family;
    pool_info.flags = vk::CommandPoolCreateFlagBits::eResetCommandBuffer;

    try {
        m_command_pool = device.createCommandPool(pool_info);
    } catch (const vk::SystemError& e) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Failed to create command pool: {}", e.what());
        return false;
    }

    pool_info.queueFamilyIndex = compute_queue_family;
    try {
        m_compute_command_pool = device.createCommandPool(pool_info);
    } catch (const vk::SystemError& e) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Failed to create compute command pool: {}", e.what());
        return false;
    }

    MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "Command manager initialized");

    return true;
}

void VKCommandManager::cleanup()
{
    if (m_compute_command_pool) {
        if (!m_compute_allocated_buffers.empty()) {
            m_device.freeCommandBuffers(m_compute_command_pool, m_compute_allocated_buffers);
            m_compute_allocated_buffers.clear();
        }
        m_device.destroyCommandPool(m_compute_command_pool);
        m_compute_command_pool = nullptr;
    }

    if (m_command_pool) {
        if (!m_allocated_buffers.empty()) {
            m_device.freeCommandBuffers(m_command_pool, m_allocated_buffers);
            m_allocated_buffers.clear();
        }

        m_device.destroyCommandPool(m_command_pool);
        m_command_pool = nullptr;
    }
}

vk::CommandBuffer VKCommandManager::allocate_command_buffer(vk::CommandBufferLevel level)
{
    vk::CommandBufferAllocateInfo alloc_info {};
    alloc_info.commandPool = m_command_pool;
    alloc_info.level = level;
    alloc_info.commandBufferCount = 1;

    vk::CommandBuffer cmd_buffer;
    try {
        auto result = m_device.allocateCommandBuffers(&alloc_info, &cmd_buffer);
        if (result != vk::Result::eSuccess) {
            MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "Failed to allocate command buffer");
            return nullptr;
        }
        m_allocated_buffers.push_back(cmd_buffer);
    } catch (const vk::SystemError& e) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Failed to allocate command buffer: {}", e.what());
        return nullptr;
    }

    MF_RT_TRACE(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "Allocated {} command buffer",
        level == vk::CommandBufferLevel::ePrimary ? "primary" : "secondary");

    return cmd_buffer;
}

void VKCommandManager::free_command_buffer(vk::CommandBuffer command_buffer)
{
    if (!command_buffer) {
        return;
    }

    auto it = std::ranges::find(m_allocated_buffers, command_buffer);
    if (it != m_allocated_buffers.end()) {
        m_device.freeCommandBuffers(m_command_pool, 1, &command_buffer);
        m_allocated_buffers.erase(it);
    }

    auto it2 = std::ranges::find(m_compute_allocated_buffers, command_buffer);
    if (it2 != m_compute_allocated_buffers.end()) {
        m_device.freeCommandBuffers(m_compute_command_pool, 1, &command_buffer);
        m_compute_allocated_buffers.erase(it2);
    }
}

vk::CommandBuffer VKCommandManager::begin_single_time_commands()
{
    vk::CommandBuffer command_buffer = allocate_command_buffer();

    vk::CommandBufferBeginInfo begin_info {};
    begin_info.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit;

    command_buffer.begin(begin_info);

    return command_buffer;
}

vk::CommandBuffer VKCommandManager::begin_single_time_commands_compute()
{
    vk::CommandBufferAllocateInfo alloc_info {};
    alloc_info.commandPool = m_compute_command_pool;
    alloc_info.level = vk::CommandBufferLevel::ePrimary;
    alloc_info.commandBufferCount = 1;

    vk::CommandBuffer cmd = m_device.allocateCommandBuffers(alloc_info).front();
    m_compute_allocated_buffers.push_back(cmd);

    vk::CommandBufferBeginInfo begin_info {};
    begin_info.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit;
    cmd.begin(begin_info);
    return cmd;
}

void VKCommandManager::end_single_time_commands(vk::CommandBuffer command_buffer, vk::Queue queue)
{
    command_buffer.end();

    vk::SubmitInfo submit_info {};
    submit_info.commandBufferCount = 1;
    submit_info.pCommandBuffers = &command_buffer;

    if (auto result = queue.submit(1, &submit_info, nullptr); result != vk::Result::eSuccess) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Failed to submit single time command buffer: {}", vk::to_string(result));
    }
    queue.waitIdle();

    free_command_buffer(command_buffer);
}

void VKCommandManager::record_deferred_commands(
    const std::function<void(vk::CommandBuffer)>& recorder)
{
    if (!m_deferred_cmd)
        m_deferred_cmd = begin_single_time_commands();
    recorder(m_deferred_cmd);
    m_deferred_pending = true;
}

void VKCommandManager::reset_deferred()
{
    m_deferred_cmd = nullptr;
    m_deferred_pending = false;
}

void VKCommandManager::reset_pool()
{
    m_device.resetCommandPool(m_command_pool, vk::CommandPoolResetFlags {});
}

} // namespace MayaFlux::Core
