#include "XCEngine/RHI/Vulkan/VulkanCommandList.h"

#include "XCEngine/RHI/Vulkan/VulkanDevice.h"
#include "XCEngine/RHI/Vulkan/VulkanResourceView.h"
#include "XCEngine/RHI/Vulkan/VulkanTexture.h"

namespace XCEngine {
namespace RHI {

namespace {

VkImageLayout ToVulkanImageLayout(ResourceStates state) {
    switch (state) {
        case ResourceStates::RenderTarget:
            return VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        case ResourceStates::CopySrc:
            return VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        case ResourceStates::CopyDst:
            return VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        case ResourceStates::Present:
            return VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
        case ResourceStates::PixelShaderResource:
            return VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        case ResourceStates::DepthWrite:
            return VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
        case ResourceStates::Common:
        default:
            return VK_IMAGE_LAYOUT_UNDEFINED;
    }
}

VkAccessFlags ToVulkanAccessMask(ResourceStates state) {
    switch (state) {
        case ResourceStates::RenderTarget:
            return VK_ACCESS_TRANSFER_WRITE_BIT;
        case ResourceStates::CopySrc:
            return VK_ACCESS_TRANSFER_READ_BIT;
        case ResourceStates::CopyDst:
            return VK_ACCESS_TRANSFER_WRITE_BIT;
        case ResourceStates::PixelShaderResource:
            return VK_ACCESS_SHADER_READ_BIT;
        case ResourceStates::DepthWrite:
            return VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
        case ResourceStates::Present:
        case ResourceStates::Common:
        default:
            return 0;
    }
}

VkPipelineStageFlags ToVulkanStageMask(ResourceStates state) {
    switch (state) {
        case ResourceStates::RenderTarget:
        case ResourceStates::CopySrc:
        case ResourceStates::CopyDst:
            return VK_PIPELINE_STAGE_TRANSFER_BIT;
        case ResourceStates::PixelShaderResource:
            return VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
        case ResourceStates::DepthWrite:
            return VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT | VK_PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT;
        case ResourceStates::Present:
            return VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
        case ResourceStates::Common:
        default:
            return VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
    }
}

} // namespace

VulkanCommandList::~VulkanCommandList() {
    Shutdown();
}

bool VulkanCommandList::Initialize(VulkanDevice* device) {
    if (device == nullptr || device->GetDevice() == VK_NULL_HANDLE) {
        return false;
    }

    m_device = device;

    VkCommandPoolCreateInfo poolInfo = {};
    poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    poolInfo.queueFamilyIndex = device->GetGraphicsQueueFamilyIndex();
    if (vkCreateCommandPool(device->GetDevice(), &poolInfo, nullptr, &m_commandPool) != VK_SUCCESS) {
        return false;
    }

    VkCommandBufferAllocateInfo allocateInfo = {};
    allocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocateInfo.commandPool = m_commandPool;
    allocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocateInfo.commandBufferCount = 1;
    return vkAllocateCommandBuffers(device->GetDevice(), &allocateInfo, &m_commandBuffer) == VK_SUCCESS;
}

void VulkanCommandList::Shutdown() {
    if (m_commandPool != VK_NULL_HANDLE && m_device != nullptr) {
        if (m_commandBuffer != VK_NULL_HANDLE) {
            vkFreeCommandBuffers(m_device->GetDevice(), m_commandPool, 1, &m_commandBuffer);
        }
        vkDestroyCommandPool(m_device->GetDevice(), m_commandPool, nullptr);
    }

    m_commandBuffer = VK_NULL_HANDLE;
    m_commandPool = VK_NULL_HANDLE;
    m_currentColorTarget = nullptr;
    m_currentDepthTarget = nullptr;
    m_device = nullptr;
}

void VulkanCommandList::Reset() {
    if (m_device == nullptr || m_commandPool == VK_NULL_HANDLE || m_commandBuffer == VK_NULL_HANDLE) {
        return;
    }

    vkResetCommandPool(m_device->GetDevice(), m_commandPool, 0);

    VkCommandBufferBeginInfo beginInfo = {};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    vkBeginCommandBuffer(m_commandBuffer, &beginInfo);
}

void VulkanCommandList::Close() {
    if (m_commandBuffer != VK_NULL_HANDLE) {
        vkEndCommandBuffer(m_commandBuffer);
    }
}

void VulkanCommandList::TransitionTexture(VulkanTexture* texture, ResourceStates newState) {
    if (texture == nullptr || m_commandBuffer == VK_NULL_HANDLE) {
        return;
    }

    VkImageMemoryBarrier barrier = {};
    barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    barrier.oldLayout = ToVulkanImageLayout(texture->GetState());
    barrier.newLayout = ToVulkanImageLayout(newState);
    barrier.srcAccessMask = ToVulkanAccessMask(texture->GetState());
    barrier.dstAccessMask = ToVulkanAccessMask(newState);
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.image = texture->GetImage();
    barrier.subresourceRange.aspectMask = GetImageAspectMask(texture->GetFormat());
    barrier.subresourceRange.baseMipLevel = 0;
    barrier.subresourceRange.levelCount = 1;
    barrier.subresourceRange.baseArrayLayer = 0;
    barrier.subresourceRange.layerCount = 1;

    vkCmdPipelineBarrier(
        m_commandBuffer,
        ToVulkanStageMask(texture->GetState()),
        ToVulkanStageMask(newState),
        0,
        0, nullptr,
        0, nullptr,
        1, &barrier);

    texture->SetState(newState);
}

void VulkanCommandList::TransitionBarrier(RHIResourceView* resource, ResourceStates stateBefore, ResourceStates stateAfter) {
    (void)stateBefore;
    auto* view = static_cast<VulkanResourceView*>(resource);
    if (view != nullptr) {
        TransitionTexture(view->GetTexture(), stateAfter);
    }
}

void VulkanCommandList::BeginRenderPass(class RHIRenderPass* renderPass, class RHIFramebuffer* framebuffer, const Rect& renderArea, uint32_t clearValueCount, const ClearValue* clearValues) {
    (void)renderPass;
    (void)framebuffer;
    (void)renderArea;
    (void)clearValueCount;
    (void)clearValues;
}

void VulkanCommandList::EndRenderPass() {
}

void VulkanCommandList::SetShader(RHIShader* shader) {
    (void)shader;
}

void VulkanCommandList::SetPipelineState(RHIPipelineState* pso) {
    (void)pso;
}

void VulkanCommandList::SetGraphicsDescriptorSets(uint32_t firstSet, uint32_t count, RHIDescriptorSet** descriptorSets, RHIPipelineLayout* pipelineLayout) {
    (void)firstSet;
    (void)count;
    (void)descriptorSets;
    (void)pipelineLayout;
}

void VulkanCommandList::SetComputeDescriptorSets(uint32_t firstSet, uint32_t count, RHIDescriptorSet** descriptorSets, RHIPipelineLayout* pipelineLayout) {
    (void)firstSet;
    (void)count;
    (void)descriptorSets;
    (void)pipelineLayout;
}

void VulkanCommandList::SetPrimitiveTopology(PrimitiveTopology topology) {
    (void)topology;
}

void VulkanCommandList::SetViewport(const Viewport& viewport) {
    (void)viewport;
}

void VulkanCommandList::SetViewports(uint32_t count, const Viewport* viewports) {
    (void)count;
    (void)viewports;
}

void VulkanCommandList::SetScissorRect(const Rect& rect) {
    (void)rect;
}

void VulkanCommandList::SetScissorRects(uint32_t count, const Rect* rects) {
    (void)count;
    (void)rects;
}

void VulkanCommandList::SetRenderTargets(uint32_t count, RHIResourceView** renderTargets, RHIResourceView* depthStencil) {
    m_currentColorTarget = (count > 0 && renderTargets != nullptr) ? renderTargets[0] : nullptr;
    m_currentDepthTarget = depthStencil;
}

void VulkanCommandList::SetStencilRef(uint8_t ref) {
    (void)ref;
}

void VulkanCommandList::SetBlendFactor(const float factor[4]) {
    (void)factor;
}

void VulkanCommandList::SetVertexBuffers(uint32_t startSlot, uint32_t count, RHIResourceView** buffers, const uint64_t* offsets, const uint32_t* strides) {
    (void)startSlot;
    (void)count;
    (void)buffers;
    (void)offsets;
    (void)strides;
}

void VulkanCommandList::SetIndexBuffer(RHIResourceView* buffer, uint64_t offset) {
    (void)buffer;
    (void)offset;
}

void VulkanCommandList::Draw(uint32_t vertexCount, uint32_t instanceCount, uint32_t startVertex, uint32_t startInstance) {
    (void)vertexCount;
    (void)instanceCount;
    (void)startVertex;
    (void)startInstance;
}

void VulkanCommandList::DrawIndexed(uint32_t indexCount, uint32_t instanceCount, uint32_t startIndex, int32_t baseVertex, uint32_t startInstance) {
    (void)indexCount;
    (void)instanceCount;
    (void)startIndex;
    (void)baseVertex;
    (void)startInstance;
}

void VulkanCommandList::Clear(float r, float g, float b, float a, uint32_t buffers) {
    (void)buffers;

    auto* colorView = static_cast<VulkanResourceView*>(m_currentColorTarget);
    if (colorView == nullptr || colorView->GetTexture() == nullptr) {
        return;
    }

    VulkanTexture* texture = colorView->GetTexture();
    TransitionTexture(texture, ResourceStates::RenderTarget);

    VkClearColorValue clearColor = {};
    clearColor.float32[0] = r;
    clearColor.float32[1] = g;
    clearColor.float32[2] = b;
    clearColor.float32[3] = a;

    VkImageSubresourceRange range = {};
    range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    range.baseMipLevel = 0;
    range.levelCount = 1;
    range.baseArrayLayer = 0;
    range.layerCount = 1;

    vkCmdClearColorImage(
        m_commandBuffer,
        texture->GetImage(),
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        &clearColor,
        1,
        &range);

    TransitionTexture(texture, ResourceStates::Present);
}

void VulkanCommandList::ClearRenderTarget(RHIResourceView* renderTarget, const float color[4]) {
    m_currentColorTarget = renderTarget;
    Clear(color[0], color[1], color[2], color[3], 1);
}

void VulkanCommandList::ClearDepthStencil(RHIResourceView* depthStencil, float depth, uint8_t stencil) {
    (void)depthStencil;
    (void)depth;
    (void)stencil;
}

void VulkanCommandList::CopyResource(RHIResourceView* dst, RHIResourceView* src) {
    (void)dst;
    (void)src;
}

void VulkanCommandList::Dispatch(uint32_t x, uint32_t y, uint32_t z) {
    (void)x;
    (void)y;
    (void)z;
}

} // namespace RHI
} // namespace XCEngine