#include <destrum/Graphics/Pipelines/SkinningPipeline.h>

#include "destrum/FS/AssetFS.h"
#include "../../../include/destrum/Graphics/Caches/MeshCache.h"
#include "destrum/Graphics/MeshDrawCommand.h"
#include "destrum/Graphics/Util.h"

#include <array>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <stdexcept>

#include "volk.h"

void SkinningPipeline::init(GfxDevice& gfxDevice) {
    try {
    const auto& device = gfxDevice.getDevice();

    const auto pushConstant = VkPushConstantRange{
        .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
        .offset = 0,
        .size = sizeof(PushConstants),
    };

    const auto pushConstants = std::array{pushConstant};

    VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.pushConstantRangeCount = 1;
    pipelineLayoutInfo.pPushConstantRanges = pushConstants.data();

    if (vkCreatePipelineLayout(device, &pipelineLayoutInfo, nullptr, &m_pipelineLayout) != VK_SUCCESS) {
        throw std::runtime_error("Could not make pipeline layout");
    }

    ComputePipelineConfigInfo pipelineConfig{};
    pipelineConfig.name = "skinning compute pipeline";
    pipelineConfig.pipelineLayout = m_pipelineLayout;

    const auto SkinningShaderPath = AssetFS::GetInstance().GetCookedPathForFile("engine://shaders/skinning.comp");

    skinningPipeline = std::make_unique<ComputePipeline>(
        gfxDevice,
        SkinningShaderPath.generic_string(),
        pipelineConfig
    );

        for (std::size_t i = 0; i < FRAMES_IN_FLIGHT; ++i) {
        auto& jointMatricesBuffer = framesData[i].jointMatricesBuffer;
        jointMatricesBuffer.capacity = MAX_JOINT_MATRICES;
        jointMatricesBuffer.buffer = gfxDevice.createBuffer(
            MAX_JOINT_MATRICES * sizeof(glm::mat4),
            VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
            VMA_MEMORY_USAGE_CPU_TO_GPU);

        const std::string bufferName = "joint matrices (frame " + std::to_string(i) + ")";
        vkutil::addDebugLabel(device, jointMatricesBuffer.buffer.buffer, bufferName.c_str());
        }
    } catch (...) {
        cleanup(gfxDevice);
        throw;
    }
}

void SkinningPipeline::cleanup(GfxDevice& gfxDevice) {
    for (auto& frame : framesData) {
        gfxDevice.destroyBuffer(frame.jointMatricesBuffer.buffer);
        frame.jointMatricesBuffer.buffer = {};
        frame.jointMatricesBuffer.size = 0;
    }

    if (m_pipelineLayout != VK_NULL_HANDLE && gfxDevice.getDevice() != VK_NULL_HANDLE) {
        vkDestroyPipelineLayout(gfxDevice.getDevice(), m_pipelineLayout, nullptr);
    }
    m_pipelineLayout = VK_NULL_HANDLE;
    skinningPipeline.reset();
}

void SkinningPipeline::doSkinning(VkCommandBuffer cmd,
                                  std::size_t frameIndex,
                                  const MeshCache& meshCache,
                                  const MeshDrawCommand& dc) {
    skinningPipeline->bind(cmd);

    const auto& mesh = meshCache.getMesh(dc.meshId);
    assert(mesh.hasSkeleton);
    assert(dc.skinnedMesh);

    const std::string& meshKey = meshCache.getMeshKey(dc.meshId);
    vkutil::ScopedCmdLabel label(cmd, meshKey.empty() ? "Skin Mesh" : meshKey.c_str(), vkutil::labelColors::compute);

    const auto cs = PushConstants{
        .jointMatricesBuffer = getCurrentFrameData(frameIndex).jointMatricesBuffer.buffer.address,
        .jointMatricesStartIndex = static_cast<std::uint32_t>(dc.jointMatricesStartIndex),
        .numVertices = mesh.numVertices,
        .inputBuffer = mesh.vertexBuffer.address,
        .skinningData = mesh.skinningDataBuffer.address,
        .outputBuffer = dc.skinnedMesh->skinnedVertexBuffer.address,
    };

    vkCmdPushConstants(cmd,
                       m_pipelineLayout,
                       VK_SHADER_STAGE_COMPUTE_BIT,
                       0,
                       sizeof(PushConstants),
                       &cs);

    constexpr std::uint32_t workgroupSize = 256;
    const auto groupSizeX = static_cast<std::uint32_t>((mesh.numVertices + workgroupSize - 1) / workgroupSize);
    if (groupSizeX == 0) {
        return;
    }

    vkCmdDispatch(cmd, groupSizeX, 1, 1);

    const VkBufferMemoryBarrier2 skinnedVertexBarrier{
        .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2,
        .srcStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT,
        .srcAccessMask = VK_ACCESS_2_SHADER_STORAGE_WRITE_BIT,
        .dstStageMask = VK_PIPELINE_STAGE_2_VERTEX_SHADER_BIT,
        .dstAccessMask = VK_ACCESS_2_SHADER_STORAGE_READ_BIT,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .buffer = dc.skinnedMesh->skinnedVertexBuffer.buffer,
        .offset = 0,
        .size = VK_WHOLE_SIZE,
    };

    const VkDependencyInfo dependencyInfo{
        .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
        .bufferMemoryBarrierCount = 1,
        .pBufferMemoryBarriers = &skinnedVertexBarrier,
    };
    vkCmdPipelineBarrier2(cmd, &dependencyInfo);
}

void SkinningPipeline::beginDrawing(std::size_t frameIndex) {
    getCurrentFrameData(frameIndex).jointMatricesBuffer.clear();
}

void SkinningPipeline::flushCurrentFrame(
    VkCommandBuffer cmd,
    GfxDevice& gfxDevice,
    std::size_t frameIndex)
{
    auto& buffer = getCurrentFrameData(frameIndex).jointMatricesBuffer.buffer;
    gfxDevice.getMemoryManager().flushAllocation(buffer, 0, VK_WHOLE_SIZE);
    vkutil::bufferHostWriteToShaderReadBarrier(
        cmd,
        buffer.buffer,
        0,
        static_cast<VkDeviceSize>(MAX_JOINT_MATRICES * sizeof(glm::mat4)));
}

std::size_t SkinningPipeline::appendJointMatrices(std::span<const glm::mat4> jointMatrices,
                                                  std::size_t frameIndex) {
    auto& jointMatricesBuffer = getCurrentFrameData(frameIndex).jointMatricesBuffer;
    const auto startIndex = jointMatricesBuffer.size;
    jointMatricesBuffer.append(jointMatrices);
    return startIndex;
}

SkinningPipeline::PerFrameData& SkinningPipeline::getCurrentFrameData(std::size_t frameIndex) {
    return framesData[frameIndex % FRAMES_IN_FLIGHT];
}