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

#include <array>
#include <cassert>
#include <stdexcept>

#include <destrum/FS/AssetFS.h>
#include <destrum/Graphics/GfxDevice.h>
#include <destrum/Graphics/RenderResources.h>
#include <destrum/Graphics/Caches/MeshCache.h>
#include <destrum/Graphics/Util.h>

#include "volk.h"
#include "spdlog/spdlog.h"

MeshPipeline::MeshPipeline() = default;

MeshPipeline::~MeshPipeline() = default;

void MeshPipeline::init(
    GfxDevice& gfxDevice,
    RenderResources& resources,
    VkFormat drawImageFormat,
    VkFormat depthImageFormat,
    VkFormat selectionImageFormat)
{
    try
    {
        const auto vertexShader =
            AssetFS::GetInstance().GetCookedPathForFile("engine://shaders/mesh.vert");

        const auto fragShader =
            AssetFS::GetInstance().GetCookedPathForFile("engine://shaders/mesh.frag");

        const auto selectionShader =
            AssetFS::GetInstance().GetCookedPathForFile("engine://shaders/selection.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{
            resources.getBindlessDescSetLayout()
        };

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

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

        {
            PipelineConfigInfo pipelineConfig{};
            Pipeline::DefaultPipelineConfigInfo(pipelineConfig);

            pipelineConfig.name = "Mesh Depth Prepass Pipeline";
            pipelineConfig.pipelineLayout = m_pipelineLayout;

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

            pipelineConfig.colorAttachments = {};
            pipelineConfig.colorBlendInfo.attachmentCount = 0;
            pipelineConfig.depthAttachment = depthImageFormat;

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

        {
            PipelineConfigInfo pipelineConfig{};
            Pipeline::DefaultPipelineConfigInfo(pipelineConfig);

            pipelineConfig.name = "Mesh Selection Prepass Pipeline";
            pipelineConfig.pipelineLayout = m_pipelineLayout;

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

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

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

        {
            PipelineConfigInfo pipelineConfig{};
            Pipeline::DefaultPipelineConfigInfo(pipelineConfig);

            pipelineConfig.name = "Mesh Pipeline";
            pipelineConfig.pipelineLayout = m_pipelineLayout;

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

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

            // depth was already written by the pre-pass, only shade the visible surface
            pipelineConfig.depthStencilInfo.depthWriteEnable = VK_FALSE;
            pipelineConfig.depthStencilInfo.depthCompareOp = VK_COMPARE_OP_LESS_OR_EQUAL;

            m_pipeline = std::make_unique<Pipeline>(
                gfxDevice,
                vertexShader.string(),
                fragShader.string(),
                pipelineConfig
            );
        }
    }
    catch (...)
    {
        cleanup(gfxDevice.getDevice());
        throw;
    }
}

void MeshPipeline::drawDepthPrepass(
    VkCommandBuffer cmd,
    VkExtent2D renderExtent,
    const RenderResources& resources,
    const GPUBuffer& sceneDataBuffer,
    const IndirectDrawList& drawList,
    const IndirectDrawList::View& view,
    bool writeObjectIds)
{
    const auto& pipeline = writeObjectIds ? m_selectionPrepassPipeline : m_depthPrepassPipeline;
    if (!pipeline)
    {
        return;
    }

    drawMeshes(
        cmd,
        *pipeline,
        renderExtent,
        resources,
        sceneDataBuffer,
        drawList,
        view);
}

void MeshPipeline::draw(
    VkCommandBuffer cmd,
    VkExtent2D renderExtent,
    const RenderResources& resources,
    const GPUBuffer& sceneDataBuffer,
    const IndirectDrawList& drawList,
    const IndirectDrawList::View& view)
{
    if (!m_pipeline)
    {
        return;
    }

    drawMeshes(
        cmd,
        *m_pipeline,
        renderExtent,
        resources,
        sceneDataBuffer,
        drawList,
        view);
}

void MeshPipeline::drawMeshes(
    VkCommandBuffer cmd,
    const Pipeline& pipeline,
    VkExtent2D renderExtent,
    const RenderResources& resources,
    const GPUBuffer& sceneDataBuffer,
    const IndirectDrawList& drawList,
    const IndirectDrawList::View& view)
{
    pipeline.bind(cmd);

    resources.bindBindlessDescSet(cmd, m_pipelineLayout);

    const auto viewport = VkViewport{
        .x = 0.f,
        .y = 0.f,
        .width = static_cast<float>(renderExtent.width),
        .height = static_cast<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);

    vkCmdSetPolygonModeEXT(
        cmd,
        m_renderWireframe ? VK_POLYGON_MODE_LINE : VK_POLYGON_MODE_FILL
    );

    // The same for every draw, per-draw data (transform, material, vertex buffer)
    // lives in the instance buffer and is looked up with gl_InstanceIndex
    const auto pushConstants = PushConstants{
        .sceneDataBuffer = sceneDataBuffer.address,
        .instanceBuffer = drawList.getInstanceBufferAddress(),
    };

    vkCmdPushConstants(
        cmd,
        m_pipelineLayout,
        VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT,
        0,
        sizeof(PushConstants),
        &pushConstants
    );

    vkutil::cmdBeginLabelFormatted(
        cmd,
        vkutil::labelColors::draw,
        "%u indirect draws, %u instances",
        view.commandCount,
        view.instanceCount);

    drawList.draw(cmd, resources.meshes().getIndexBuffer(), view);

    vkutil::cmdEndLabel(cmd);
}

void MeshPipeline::cleanup(VkDevice device)
{
    m_pipeline.reset();
    m_depthPrepassPipeline.reset();
    m_selectionPrepassPipeline.reset();

    if (m_pipelineLayout != VK_NULL_HANDLE && device != VK_NULL_HANDLE)
    {
        vkDestroyPipelineLayout(device, m_pipelineLayout, nullptr);
    }
    m_pipelineLayout = VK_NULL_HANDLE;
}