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

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

#include <glm/gtc/matrix_transform.hpp>

#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"

void ShadowMapPipeline::init(GfxDevice& gfxDevice, RenderResources& resources)
{
    try
    {
        const auto vertexShader =
            AssetFS::GetInstance().GetCookedPathForFile("engine://shaders/shadow.vert");

        const auto bufferRange = VkPushConstantRange{
            .stageFlags = VK_SHADER_STAGE_VERTEX_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 shadow map pipeline layout");
        }

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

        pipelineConfig.name = "Shadow Map Pipeline";
        pipelineConfig.pipelineLayout = m_pipelineLayout;

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

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

        // the light projection isn't Y flipped like the camera's, so the winding is mirrored
        pipelineConfig.rasterizationInfo.cullMode = VK_CULL_MODE_NONE;
        pipelineConfig.rasterizationInfo.depthBiasEnable = VK_TRUE;
        pipelineConfig.rasterizationInfo.depthBiasConstantFactor = 0.0f;
        pipelineConfig.rasterizationInfo.depthBiasSlopeFactor = 2.0f;

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

        const auto createInfo = vkutil::CreateImageInfo{
            .format = shadowMapFormat,
            .usage = VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
            .extent = VkExtent3D{SHADOW_MAP_SIZE, SHADOW_MAP_SIZE, 1},
        };

        shadowMapId = resources.createImage(gfxDevice, createInfo, "shadow map", nullptr, shadowMapId);
        spdlog::info("Created {0}x{0} shadow map with id {1}", SHADOW_MAP_SIZE, shadowMapId);
    }
    catch (...)
    {
        cleanup(gfxDevice.getDevice());
        throw;
    }
}

void ShadowMapPipeline::cleanup(VkDevice device)
{
    m_pipeline.reset();

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

    shadowMapId = NULL_IMAGE_ID;
}

glm::mat4 ShadowMapPipeline::calculateLightViewProj(const glm::vec3& lightDirection, const glm::vec3& center) const
{
    const glm::vec3 direction = glm::normalize(lightDirection);
    const glm::vec3 up = std::abs(direction.y) > 0.99f ? glm::vec3{0.0f, 0.0f, 1.0f} : glm::vec3{0.0f, 1.0f, 0.0f};

    const glm::mat4 lightView = glm::lookAt(glm::vec3{0.0f}, direction, up);

    const float texelSize = (2.0f * shadowDistance) / static_cast<float>(SHADOW_MAP_SIZE);
    glm::vec3 lightSpaceCenter = glm::vec3(lightView * glm::vec4(center, 1.0f));
    lightSpaceCenter.x = std::floor(lightSpaceCenter.x / texelSize) * texelSize;
    lightSpaceCenter.y = std::floor(lightSpaceCenter.y / texelSize) * texelSize;

    const glm::mat4 lightProj = glm::orthoRH_ZO(
        lightSpaceCenter.x - shadowDistance,
        lightSpaceCenter.x + shadowDistance,
        lightSpaceCenter.y - shadowDistance,
        lightSpaceCenter.y + shadowDistance,
        -lightSpaceCenter.z - shadowDepthRange,
        -lightSpaceCenter.z + shadowDepthRange);

    return lightProj * lightView;
}

float ShadowMapPipeline::getNormalOffset() const
{
    const float texelSize = (2.0f * shadowDistance) / static_cast<float>(SHADOW_MAP_SIZE);
    return texelSize * 1.5f;
}

void ShadowMapPipeline::draw(
    VkCommandBuffer cmd,
    const RenderResources& resources,
    const GPUBuffer& sceneDataBuffer,
    const std::vector<MeshDrawCommand>& drawCommands,
    const std::vector<std::size_t>& sortedDrawCommands)
{
    if (!m_pipeline || shadowMapId == NULL_IMAGE_ID)
    {
        return;
    }

    const auto& shadowMap = resources.getImage(shadowMapId);

    vkutil::transitionImage(
        cmd,
        shadowMap.image,
        shadowMap.layout,
        VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL);
    shadowMap.setLayout(VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL);

    auto renderInfo = vkutil::createRenderingInfo({
        .renderExtent = shadowMap.getExtent2D(),
        .depthImageView = shadowMap.imageView,
        .depthImageClearValue = 1.f,
    });
    renderInfo.renderingInfo.pDepthAttachment = &renderInfo.depthAttachment;

    vkCmdBeginRendering(cmd, &renderInfo.renderingInfo);

    m_pipeline->bind(cmd);
    resources.bindBindlessDescSet(cmd, m_pipelineLayout);

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

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

    vkCmdSetPolygonModeEXT(cmd, VK_POLYGON_MODE_FILL);

    const MeshCache& meshCache = resources.meshes();
    MeshID prevMeshId = NULL_MESH_ID;

    const auto drawOne = [&](const MeshDrawCommand& dc)
    {
        if (!dc.castShadow)
        {
            return;
        }

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

        const std::string& meshKey = meshCache.getMeshKey(dc.meshId);
        vkutil::cmdBeginLabelFormatted(
            cmd,
            vkutil::labelColors::draw,
            "%s [mesh %zu, material %u%s]",
            meshKey.empty() ? "Mesh" : meshKey.c_str(),
            dc.meshId,
            dc.materialId,
            dc.skinnedMesh != nullptr ? ", skinned" : "");

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

        const auto pushConstants = PushConstants{
            .transform = dc.transformMatrix,
            .sceneDataBuffer = sceneDataBuffer.address,
            .vertexBuffer =
                dc.skinnedMesh != nullptr
                    ? dc.skinnedMesh->skinnedVertexBuffer.address
                    : mesh.vertexBuffer.address,
            .materialId = dc.materialId,
            .padding = 0,
        };

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

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

        vkutil::cmdEndLabel(cmd);
    };

    if (!sortedDrawCommands.empty())
    {
        for (std::size_t drawIndex : sortedDrawCommands)
        {
            assert(drawIndex < drawCommands.size());
            drawOne(drawCommands[drawIndex]);
        }
    }
    else
    {
        for (const MeshDrawCommand& dc : drawCommands)
        {
            drawOne(dc);
        }
    }

    vkCmdEndRendering(cmd);

    vkutil::transitionImage(
        cmd,
        shadowMap.image,
        shadowMap.layout,
        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    shadowMap.setLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
}