#include <destrum/Graphics/Pipelines/MeshPipeline.h>
#include <destrum/FS/AssetFS.h>

MeshPipeline::MeshPipeline() {
}

MeshPipeline::~MeshPipeline() {
}

void MeshPipeline::init(GfxDevice& gfxDevice, VkFormat drawImageFormat, VkFormat depthImageFormat) {
    const auto& device = gfxDevice.getDevice();

    // const auto vertexShader = vkutil::loadShaderModule(AssetFS::GetInstance().GetFullPath("engine://shaders/mesh.vert.spv"), device);
    // const auto fragShader = vkutil::loadShaderModule(AssetFS::GetInstance().GetFullPath("engine://shaders/mesh.frag.spv"), device);


    const auto vertexShader = AssetFS::GetInstance().GetFullPath("engine://shaders/mesh.vert.spv");
    const auto fragShader = AssetFS::GetInstance().GetFullPath("engine://shaders/mesh.frag.spv");


    // vkutil::addDebugLabel(device, vertexShader, "mesh.vert");
    // vkutil::addDebugLabel(device, vertexShader, "mesh.frag");

    const auto bufferRange = VkPushConstantRange{
        .stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT,
        .offset = 0,
        .size = sizeof(PushConstants),
    };

    const auto pushConstantRanges = std::array{bufferRange};
    const auto layouts = std::array{gfxDevice.getBindlessDescSetLayout()};
    // m_pipelineLayout = vkutil::createPipelineLayout(device, layouts, pushConstantRanges);
    // vkutil::addDebugLabel(device, pipelineLayout, "mesh pipeline layout");

    VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.setLayoutCount = static_cast<uint32_t>(layouts.size());
    pipelineLayoutInfo.pSetLayouts = layouts.data();

    pipelineLayoutInfo.pushConstantRangeCount = 1;
    pipelineLayoutInfo.pPushConstantRanges = pushConstantRanges.data();

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

    PipelineConfigInfo pipelineConfig{};
    Pipeline::DefaultPipelineConfigInfo(pipelineConfig);
    pipelineConfig.name = "Mesh Pipeline";
    pipelineConfig.pipelineLayout = m_pipelineLayout;

    pipelineConfig.vertexAttributeDescriptions = {};
    pipelineConfig.vertexBindingDescriptions = {};

    pipelineConfig.colorAttachments = {drawImageFormat};
    pipelineConfig.depthAttachment = depthImageFormat;

    m_pipeline = std::make_unique<Pipeline>(
        gfxDevice,
        vertexShader.string(),
        fragShader.string(),
        pipelineConfig
    );
}

void MeshPipeline::draw(VkCommandBuffer cmd,
                        VkExtent2D renderExtent,
                        const GfxDevice& gfxDevice,
                        const MeshCache& meshCache,
                        const MaterialCache& materialCache,
                        const Camera& camera,
                        const GPUBuffer& sceneDataBuffer,
                        const std::vector<MeshDrawCommand>& drawCommands,
                        const std::vector<std::size_t>& sortedDrawCommands) {
    m_pipeline->bind(cmd);
    gfxDevice.bindBindlessDescSet(cmd, m_pipelineLayout);

    const auto viewport = VkViewport{
        .x = 0,
        .y = 0,
        .width = (float)renderExtent.width,
        .height = (float)renderExtent.height,
        .minDepth = 0.f,
        .maxDepth = 1.f,
    };
    vkCmdSetViewport(cmd, 0, 1, &viewport);

    const auto scissor = VkRect2D{
        .offset = {},
        .extent = renderExtent,
    };
    vkCmdSetScissor(cmd, 0, 1, &scissor);

    auto prevMeshId = NULL_MESH_ID;

    // const auto frustum = edge::createFrustumFromCamera(camera);

    for (const auto& dcIdx : drawCommands) {
        // const auto& dc = drawCommands[dcIdx];
        const auto& dc = dcIdx;

        // if (!edge::isInFrustum(frustum, dc.worldBoundingSphere)) {
            // continue;
        // }

        const auto& mesh = meshCache.getMesh(dc.meshId);
        if (dc.meshId != prevMeshId) {
            prevMeshId = dc.meshId;
            vkCmdBindIndexBuffer(cmd, mesh.indexBuffer.buffer, 0, VK_INDEX_TYPE_UINT32);
        }

        assert(dc.materialId != NULL_MATERIAL_ID);
        const auto pushConstants = PushConstants{
            .transform = dc.transformMatrix,
            .sceneDataBuffer = sceneDataBuffer.address,
            .vertexBuffer = mesh.vertexBuffer.address,
            .materialId = dc.materialId,
        };
        vkCmdPushConstants(
            cmd,
            m_pipelineLayout,
            VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT,
            0,
            sizeof(PushConstants),
            &pushConstants);

        vkCmdDrawIndexed(cmd, mesh.numIndices, 1, 0, 0, 0);
    }
}