// Copyright (c) Wojciech Figat. All rights reserved.

#if GRAPHICS_API_VULKAN

#include "CmdBufferVulkan.h"
#include "RenderToolsVulkan.h"
#include "QueueVulkan.h"
#include "GPUContextVulkan.h"
#if VULKAN_USE_TIMER_QUERIES
#include "GPUTimerQueryVulkan.h"
#endif
#include "DescriptorSetVulkan.h"
#include "Engine/Engine/Engine.h"
#include "Engine/Profiler/ProfilerCPU.h"

void CmdBufferVulkan::AddWaitSemaphore(VkPipelineStageFlags waitFlags, SemaphoreVulkan* waitSemaphore)
{
    _waitFlags.Add(waitFlags);
    ASSERT(!_waitSemaphores.Contains(waitSemaphore));
    _waitSemaphores.Add(waitSemaphore);
}

void CmdBufferVulkan::Begin()
{
    PROFILE_CPU();
    ASSERT(_state == State::ReadyForBegin);

    VkCommandBufferBeginInfo beginInfo;
    RenderToolsVulkan::ZeroStruct(beginInfo, VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO);
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    VALIDATE_VULKAN_RESULT(vkBeginCommandBuffer(_commandBuffer, &beginInfo));

    // Acquire a descriptor pool set on
    if (_descriptorPoolSetContainer == nullptr)
    {
        _descriptorPoolSetContainer = _device->DescriptorPoolsManager->AcquirePoolSetContainer();
    }

    _state = State::IsInsideBegin;

#if GPU_ALLOW_PROFILE_EVENTS
    // Reset events counter
    _eventsBegin = 0;
#endif
}

void CmdBufferVulkan::End()
{
    PROFILE_CPU();
    ASSERT(IsOutsideRenderPass());

#if GPU_ALLOW_PROFILE_EVENTS
    // End remaining events
    while (_eventsBegin--)
    {
#if VK_EXT_debug_utils
        if (vkCmdEndDebugUtilsLabelEXT)
            vkCmdEndDebugUtilsLabelEXT(GetHandle());
#endif
#if VULKAN_USE_TRACY_GPU
        tracy::EndVkZoneScope(_tracyZones.Last().Data);
        _tracyZones.RemoveLast();
#endif
    }
#endif

    VALIDATE_VULKAN_RESULT(vkEndCommandBuffer(GetHandle()));
    _state = State::HasEnded;
}

void CmdBufferVulkan::BeginRenderPass(RenderPassVulkan* renderPass, FramebufferVulkan* framebuffer, uint32 clearValueCount, VkClearValue* clearValues)
{
    ASSERT(IsOutsideRenderPass());
    VkRenderPassBeginInfo info;
    RenderToolsVulkan::ZeroStruct(info, VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO);
    info.renderPass = renderPass->Handle;
    info.framebuffer = framebuffer->Handle;
    info.renderArea.offset.x = 0;
    info.renderArea.offset.y = 0;
    info.renderArea.extent.width = framebuffer->Extent.width;
    info.renderArea.extent.height = framebuffer->Extent.height;
    info.clearValueCount = clearValueCount;
    info.pClearValues = clearValues;
    vkCmdBeginRenderPass(_commandBuffer, &info, VK_SUBPASS_CONTENTS_INLINE);
    _state = State::IsInsideRenderPass;
}

void CmdBufferVulkan::EndRenderPass()
{
    ASSERT(IsInsideRenderPass());
    vkCmdEndRenderPass(_commandBuffer);
    _state = State::IsInsideBegin;
}

#if GPU_ALLOW_PROFILE_EVENTS

void CmdBufferVulkan::BeginEvent(const Char* name, void* tracyContext)
{
    _eventsBegin++;

    char buffer[60];
    int32 bufferSize = StringUtils::Copy(buffer, name, sizeof(buffer));

#if VULKAN_USE_TRACY_GPU
    auto& zone = _tracyZones.AddOne();
    tracy::BeginVkZoneScope(zone.Data, tracyContext, GetHandle(), buffer, bufferSize);
#endif

#if VK_EXT_debug_utils
    if (vkCmdBeginDebugUtilsLabelEXT)
    {
        VkDebugUtilsLabelEXT label;
        RenderToolsVulkan::ZeroStruct(label, VK_STRUCTURE_TYPE_DEBUG_UTILS_LABEL_EXT);
        label.pLabelName = buffer;
        vkCmdBeginDebugUtilsLabelEXT(GetHandle(), &label);
    }
#endif
}

void CmdBufferVulkan::EndEvent()
{
    if (_eventsBegin == 0)
        return;
    _eventsBegin--;

#if VK_EXT_debug_utils
    if (vkCmdEndDebugUtilsLabelEXT)
        vkCmdEndDebugUtilsLabelEXT(GetHandle());
#endif

#if VULKAN_USE_TRACY_GPU
    tracy::EndVkZoneScope(_tracyZones.Last().Data);
    _tracyZones.RemoveLast();
#endif
}

#endif

void CmdBufferVulkan::Wait(float timeoutSeconds)
{
    if (!IsSubmitted())
        return;
    PROFILE_CPU();
    const bool failed = _device->FenceManager.WaitForFence(GetFence(), timeoutSeconds);
    ASSERT(!failed);
    RefreshFenceStatus();
}

void CmdBufferVulkan::RefreshFenceStatus()
{
    if (_state == State::Submitted)
    {
        PROFILE_CPU();
        if (_device->FenceManager.IsFenceSignaled(_fence))
        {
            _state = State::ReadyForBegin;
            _submittedWaitSemaphores.Clear();

            vkResetCommandBuffer(_commandBuffer, VK_COMMAND_BUFFER_RESET_RELEASE_RESOURCES_BIT);
            _device->FenceManager.ResetFence(_fence);
            _fenceSignaledCounter++;

            if (_descriptorPoolSetContainer)
            {
                _descriptorPoolSetContainer->LastFrameUsed = Engine::FrameCount;
                _descriptorPoolSetContainer = nullptr;
            }
        }
    }
    else
    {
        ASSERT(!_fence->IsSignaled);
    }
}

CmdBufferVulkan::CmdBufferVulkan(GPUDeviceVulkan* device, CmdBufferPoolVulkan* pool)
    : _device(device)
    , _commandBuffer(VK_NULL_HANDLE)
    , _state(State::ReadyForBegin)
    , _fence(nullptr)
    , _submittedFenceCounter(0)
    , _fenceSignaledCounter(0)
    , _commandBufferPool(pool)
{
    VkCommandBufferAllocateInfo createCmdBufInfo;
    RenderToolsVulkan::ZeroStruct(createCmdBufInfo, VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO);
    createCmdBufInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    createCmdBufInfo.commandBufferCount = 1;
    createCmdBufInfo.commandPool = _commandBufferPool->GetHandle();
    VALIDATE_VULKAN_RESULT(vkAllocateCommandBuffers(_device->Device, &createCmdBufInfo, &_commandBuffer));
    _fence = _device->FenceManager.AllocateFence();
}

CmdBufferVulkan::~CmdBufferVulkan()
{
    auto& fenceManager = _device->FenceManager;
    if (_state == State::Submitted)
    {
        // Wait
        fenceManager.WaitAndReleaseFence(_fence);
    }
    else
    {
        // Just free the fence, command buffer was not submitted
        fenceManager.ReleaseFence(_fence);
    }

    vkFreeCommandBuffers(_device->Device, _commandBufferPool->GetHandle(), 1, &_commandBuffer);
}

CmdBufferVulkan* CmdBufferPoolVulkan::Create()
{
    const auto cmdBuffer = New<CmdBufferVulkan>(_device, this);
    _cmdBuffers.Add(cmdBuffer);
    return cmdBuffer;
}

void CmdBufferPoolVulkan::Create(uint32 queueFamilyIndex)
{
    VkCommandPoolCreateInfo poolInfo;
    RenderToolsVulkan::ZeroStruct(poolInfo, VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO);
    poolInfo.queueFamilyIndex = queueFamilyIndex;
    // TODO: use VK_COMMAND_POOL_CREATE_TRANSIENT_BIT?
    poolInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    VALIDATE_VULKAN_RESULT(vkCreateCommandPool(_device->Device, &poolInfo, nullptr, &_handle));
}

CmdBufferPoolVulkan::CmdBufferPoolVulkan(GPUDeviceVulkan* device)
    : _device(device)
    , _handle(VK_NULL_HANDLE)
{
}

CmdBufferPoolVulkan::~CmdBufferPoolVulkan()
{
    for (int32 i = 0; i < _cmdBuffers.Count(); i++)
        Delete(_cmdBuffers[i]);
    vkDestroyCommandPool(_device->Device, _handle, nullptr);
}

void CmdBufferPoolVulkan::RefreshFenceStatus(const CmdBufferVulkan* skipCmdBuffer)
{
    for (int32 i = 0; i < _cmdBuffers.Count(); i++)
    {
        auto cmdBuffer = _cmdBuffers[i];
        if (cmdBuffer != skipCmdBuffer)
        {
            cmdBuffer->RefreshFenceStatus();
        }
    }
}

CmdBufferManagerVulkan::CmdBufferManagerVulkan(GPUDeviceVulkan* device, GPUContextVulkan* context)
    : _device(device)
    , _context(context)
    , _pool(device)
    , _queue(context->GetQueue())
    , _activeCmdBuffer(nullptr)
{
    _pool.Create(_queue->GetFamilyIndex());
}

void CmdBufferManagerVulkan::SubmitActiveCmdBuffer(SemaphoreVulkan* signalSemaphore)
{
    PROFILE_CPU();
    ASSERT(_activeCmdBuffer);
    if (!_activeCmdBuffer->IsSubmitted() && _activeCmdBuffer->HasBegun())
    {
        if (_activeCmdBuffer->IsInsideRenderPass())
            _activeCmdBuffer->EndRenderPass();

#if VULKAN_USE_TIMER_QUERIES && GPU_VULKAN_PAUSE_QUERIES
        // Pause all active timer queries
        auto queries = _activeTimerQueries.Get();
#if GPU_VULKAN_QUERY_NEW
        for (int32 i = 0; i < _activeTimerQueries.Count(); i++)
        {
            GPUQueryVulkan query;
            query.Raw = queries[i];

            // End active query to get time from start until submission
            auto pool = _device->QueryPools[query.PoolIndex];
            vkCmdWriteTimestamp(_activeCmdBuffer->GetHandle(), VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, pool->GetHandle(), query.SecondQueryIndex);
            pool->MarkQueryAsStarted(query.SecondQueryIndex);
            // TODO: somehow handle ending this query properly by stopping split query instead
            //_context->EndQuery(query.Raw);

            // TODO: reimplement timer queries pause/resume to be more exact?
        }
#else
        for (int32 i = 0; i < _activeTimerQueries.Count(); i++)
            queries->Interrupt(_activeCmdBuffer);
#endif
#endif

        _activeCmdBuffer->End();

        if (signalSemaphore)
        {
            _queue->Submit(_activeCmdBuffer, signalSemaphore->GetHandle());
        }
        else
        {
            _queue->Submit(_activeCmdBuffer);
        }
    }
    _activeCmdBuffer = nullptr;
}

void CmdBufferManagerVulkan::GetNewActiveCommandBuffer()
{
    PROFILE_CPU();
    ASSERT_LOW_LAYER(_activeCmdBuffer == nullptr)
    for (int32 i = 0; i < _pool._cmdBuffers.Count(); i++)
    {
        auto cmdBuffer = _pool._cmdBuffers.Get()[i];
        cmdBuffer->RefreshFenceStatus();
        if (cmdBuffer->GetState() == CmdBufferVulkan::State::ReadyForBegin)
        {
            _activeCmdBuffer = cmdBuffer;
            break;
        }
        else
        {
            ASSERT(cmdBuffer->GetState() == CmdBufferVulkan::State::Submitted);
        }
    }

    if (_activeCmdBuffer == nullptr)
    {
        // Always begin fresh command buffer for rendering
        _activeCmdBuffer = _pool.Create();
    }

    _activeCmdBuffer->Begin();

#if VULKAN_USE_TIMER_QUERIES && GPU_VULKAN_PAUSE_QUERIES
    // Resume any paused timer queries with the new command buffer
    auto queries = _activeTimerQueries.Get();
#if GPU_VULKAN_QUERY_NEW
    for (int32 i = 0; i < _activeTimerQueries.Count(); i++)
    {
        GPUQueryVulkan query;
        query.Raw = queries[i];
        //_activeTimerQueries.Get()[i]->Resume(_activeCmdBuffer);
    }
#else
    for (int32 i = 0; i < _activeTimerQueries.Count(); i++)
    {
        queries->Resume(_activeCmdBuffer);
    }
#endif
#endif
}

#if GPU_VULKAN_QUERY_NEW && GPU_VULKAN_PAUSE_QUERIES

void CmdBufferManagerVulkan::OnTimerQueryBegin(QueryType query)
{
    _activeTimerQueries.Add(query);
}

void CmdBufferManagerVulkan::OnTimerQueryEnd(QueryType query)
{
    _activeTimerQueries.Remove(query);
}

#endif

#endif
