#include <destrum/Graphics/Renderer.h>

#include <destrum/Graphics/Managers/LineRenderingManager.h>
#include <destrum/Graphics/Managers/Renderer2D.h>
#include <destrum/Graphics/Util.h>

#include <algorithm>
#include <numeric>

#include <glm/gtc/matrix_transform.hpp>

#include "volk.h"
#include "destrum/Util/GameState.h"
#include "spdlog/spdlog.h"

#include "tracy/TracyVulkan.hpp"

GameRenderer::GameRenderer()
{
}

void GameRenderer::init(GfxDevice& gfxDevice, RenderResources& _resources, glm::ivec2 drawImageSize)
{
    try
    {
        resources = &_resources;
        sceneDataBuffer.init(
            gfxDevice,
            VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
            sizeof(GPUSceneData),
            "scene data");

        createDrawImage(gfxDevice, drawImageSize, true);

        meshPipeline = std::make_unique<MeshPipeline>();
        meshPipeline->init(gfxDevice, _resources, drawImageFormat, depthImageFormat);

        skyboxPipeline = std::make_unique<SkyboxPipeline>();
        skyboxPipeline->init(gfxDevice, _resources, drawImageFormat, depthImageFormat);

        skinningPipeline = std::make_unique<SkinningPipeline>();
        skinningPipeline->init(gfxDevice);

        lightCullingPipeline = std::make_unique<LightCullingPipeline>();
        lightCullingPipeline->init(gfxDevice, _resources, drawImageSize);

        shadowMapPipeline = std::make_unique<ShadowMapPipeline>();
        shadowMapPipeline->init(gfxDevice, _resources);

        lineRenderingPass = std::make_unique<LineRenderingPass>();
        lineRenderingPass->init(gfxDevice, drawImageFormat);

        renderer2DPass = std::make_unique<Renderer2DPass>();
        renderer2DPass->init(gfxDevice, _resources, drawImageFormat);

        GameState::GetInstance().SetRenderer(this);
        initialized = true;
    }
    catch (...)
    {
        cleanup(gfxDevice);
        throw;
    }
}

void GameRenderer::beginDrawing(GfxDevice& gfxDevice)
{
    meshDrawCommands.clear();
    skinningPipeline->beginDrawing(gfxDevice.getCurrentFrameIndex());
    lightCullingPipeline->beginDrawing(gfxDevice.getCurrentFrameIndex());
    shadowLightDirection.reset();
    submittedLights.clear();
    spotShadows.clear();
    Renderer2D::GetInstance().Clear();
}

void GameRenderer::endDrawing()
{
    sortedMeshDrawCommands.resize(meshDrawCommands.size());
    std::iota(sortedMeshDrawCommands.begin(), sortedMeshDrawCommands.end(), 0);

    std::stable_sort(
        sortedMeshDrawCommands.begin(),
        sortedMeshDrawCommands.end(),
        [this](std::size_t left, std::size_t right)
        {
            const auto& lhs = meshDrawCommands[left];
            const auto& rhs = meshDrawCommands[right];
            if (lhs.meshId != rhs.meshId)
            {
                return lhs.meshId < rhs.meshId;
            }
            if (lhs.materialId != rhs.materialId)
            {
                return lhs.materialId < rhs.materialId;
            }
            return (lhs.skinnedMesh != nullptr) < (rhs.skinnedMesh != nullptr);
        });
}

void GameRenderer::draw(VkCommandBuffer cmd, GfxDevice& gfxDevice, const Camera& camera, const SceneData& sceneData)
{
    ZoneScopedN("GameRenderer::draw");
    TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "GameRenderer::draw");
    vkutil::ScopedCmdLabel frameLabel(cmd, "GameRenderer::draw", vkutil::labelColors::frame);

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Skinning");
        vkutil::ScopedCmdLabel label(cmd, "Skinning", vkutil::labelColors::compute);

        skinningPipeline->flushCurrentFrame(
            cmd,
            gfxDevice,
            gfxDevice.getCurrentFrameIndex());

        for (const auto& dc : meshDrawCommands)
        {
            if (!dc.skinnedMesh)
            {
                continue;
            }
            skinningPipeline->doSkinning(cmd, gfxDevice.getCurrentFrameIndex(), resources->meshes(), dc);
        }
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Lights Buffer Flush");
        vkutil::ScopedCmdLabel label(cmd, "Lights Buffer Flush", vkutil::labelColors::transfer);

        uploadLights(sceneData.camera);
        lightCullingPipeline->flushCurrentFrame(cmd, gfxDevice);
    }

    const glm::mat4 lightViewProj = shadowLightDirection
                                        ? shadowMapPipeline->calculateDirectionalViewProj(
                                            *shadowLightDirection,
                                            sceneData.camera.GetPosition())
                                        : glm::mat4{1.0f};

    std::array<GPUSpotShadow, ShadowMapPipeline::MAX_SPOT_SHADOWS> gpuSpotShadows{};
    std::ranges::copy(spotShadows, gpuSpotShadows.begin());

    std::vector<glm::mat4> spotShadowViewProjs;
    spotShadowViewProjs.reserve(spotShadows.size());
    for (const GPUSpotShadow& spotShadow : spotShadows)
    {
        spotShadowViewProjs.push_back(spotShadow.lightViewProj);
    }

    // TODO: there has to be a better way to do this
    const auto gpuSceneData = GPUSceneData{
        .view = sceneData.camera.GetViewMatrix(),
        .proj = sceneData.camera.GetProjectionMatrix(),
        .invProj = glm::inverse(sceneData.camera.GetProjectionMatrix()),
        .viewProj = sceneData.camera.GetViewProjectionMatrix(),
        .cameraPos = glm::vec4{sceneData.camera.m_position, 1.f},
        .ambientColor = {sceneData.ambientColor},
        .ambientIntensity = sceneData.ambientIntensity,
        .fogColor = {sceneData.fogColor},
        .fogDensity = sceneData.fogDensity,
        .materialsBuffer = resources->materials().getMaterialDataBufferAddress(),
        .lightsBuffer = lightCullingPipeline->getLightsBufferAddress(),
        .tileLightIndicesBuffer = lightCullingPipeline->getTileLightIndicesBufferAddress(),
        .numLights = lightCullingPipeline->getNumLights(),
        .tileCountX = lightCullingPipeline->getTileCount().x,
        .debugLightHeatmap = renderLightHeatmap ? 1u : 0u,
        .shadowMapId = static_cast<std::uint32_t>(shadowMapPipeline->getShadowMapId()),
        .shadowNormalOffset = shadowMapPipeline->getNormalOffset(),
        .padding = 0,
        .lightViewProj = lightViewProj,
        .spotShadowAtlasId = static_cast<std::uint32_t>(shadowMapPipeline->getSpotShadowAtlasId()),
        .spotShadowTilesPerRow = ShadowMapPipeline::SPOT_SHADOW_TILES_PER_ROW,
        .spotShadows = gpuSpotShadows,
    };

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Material Updates");
        vkutil::ScopedCmdLabel label(cmd, "Material Updates", vkutil::labelColors::transfer);

        recordMaterialUpdates(cmd);
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Upload Scene Data");
        vkutil::ScopedCmdLabel label(cmd, "Upload Scene Data", vkutil::labelColors::transfer);

        sceneDataBuffer.uploadNewData(
            cmd,
            gfxDevice.getCurrentFrameIndex(),
            (void*)&gpuSceneData,
            sizeof(GPUSceneData)
        );
    }

    const auto& drawImage = resources->getImage(drawImageId);
    const auto& depthImage = resources->getImage(depthImageId);

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Transition Draw Image");

        vkutil::transitionImage(
            cmd,
            drawImage.image,
            drawImage.layout,
            VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);
        drawImage.setLayout(VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Transition Depth Image");

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

    if (shadowLightDirection)
    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Shadow Map");
        vkutil::ScopedCmdLabel label(cmd, "Shadow Map", vkutil::labelColors::shadow);

        shadowMapPipeline->drawDirectionalShadow(
            cmd,
            *resources,
            meshDrawCommands,
            sortedMeshDrawCommands,
            lightViewProj);
    }

    if (!spotShadowViewProjs.empty())
    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Spot Shadows");
        vkutil::ScopedCmdLabel label(cmd, "Spot Shadows", vkutil::labelColors::shadow);

        shadowMapPipeline->drawSpotShadows(
            cmd,
            *resources,
            meshDrawCommands,
            sortedMeshDrawCommands,
            spotShadowViewProjs);
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Depth Prepass");
        vkutil::ScopedCmdLabel label(cmd, "Depth Prepass", vkutil::labelColors::depth);

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

        vkCmdBeginRendering(cmd, &depthPrepassInfo.renderingInfo);

        meshPipeline->drawDepthPrepass(
            cmd,
            drawImage.getExtent2D(),
            *resources,
            camera,
            sceneDataBuffer.getBuffer(),
            meshDrawCommands,
            sortedMeshDrawCommands);

        vkCmdEndRendering(cmd);
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Light Culling");
        vkutil::ScopedCmdLabel label(cmd, "Light Culling", vkutil::labelColors::compute);

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

        lightCullingPipeline->cullLights(
            cmd,
            *resources,
            sceneDataBuffer.getBuffer(),
            depthImageId,
            drawImage.getExtent2D());

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

    auto renderInfo = vkutil::createRenderingInfo({
        .renderExtent = drawImage.getExtent2D(),
        .colorImageView = drawImage.imageView,
        .colorImageClearValue = glm::vec4{0.f, 0.f, 0.f, 1.f},
        .depthImageView = depthImage.imageView,
    });
    renderInfo.renderingInfo.pColorAttachments = &renderInfo.colorAttachment;
    renderInfo.renderingInfo.pDepthAttachment = &renderInfo.depthAttachment;

    vkutil::cmdBeginLabel(cmd, "Forward Pass", vkutil::labelColors::opaque);

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "vkCmdBeginRendering");

        vkCmdBeginRendering(cmd, &renderInfo.renderingInfo);
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "MeshPipeline::draw");
        vkutil::ScopedCmdLabel label(cmd, "MeshPipeline::draw", vkutil::labelColors::opaque);

        meshPipeline->draw(
            cmd,
            drawImage.getExtent2D(),
            *resources,
            camera,
            sceneDataBuffer.getBuffer(),
            meshDrawCommands,
            sortedMeshDrawCommands);
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "SkyboxPipeline::draw");
        vkutil::ScopedCmdLabel label(cmd, "SkyboxPipeline::draw", vkutil::labelColors::sky);

        skyboxPipeline->draw(cmd, *resources, camera, sceneData.skyboxRotation);
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "vkCmdEndRendering");

        vkCmdEndRendering(cmd);
    }

    vkutil::cmdEndLabel(cmd);

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "LineRenderingPass::draw");
        vkutil::ScopedCmdLabel label(cmd, "LineRenderingPass::draw", vkutil::labelColors::debug);

        lineRenderingPass->draw(
            cmd,
            gfxDevice,
            drawImage,
            camera,
            LineRenderingManager::GetInstance());
    }

    {
        TracyVkZone(gfxDevice.getTracyVkCtx(), cmd, "Renderer2DPass::draw");
        vkutil::ScopedCmdLabel label(cmd, "Renderer2DPass::draw", vkutil::labelColors::ui);

        renderer2DPass->draw(
            cmd,
            gfxDevice,
            *resources,
            drawImage,
            Renderer2DPass::screenProjection(drawImage));
    }
}

void GameRenderer::draw2D(
    VkCommandBuffer cmd,
    GfxDevice& gfxDevice,
    const glm::mat4& projection)
{
    const auto& drawImage = resources->getImage(drawImageId);
    renderer2DPass->draw(cmd, gfxDevice, *resources, drawImage, projection);
}

void GameRenderer::cleanup(GfxDevice& gfxDevice)
{
    VkDevice device = gfxDevice.getDevice();

    if (skinningPipeline)
        skinningPipeline->cleanup(gfxDevice);

    if (lineRenderingPass)
        lineRenderingPass->cleanup(gfxDevice);

    if (renderer2DPass)
        renderer2DPass->cleanup(gfxDevice);

    if (lightCullingPipeline)
        lightCullingPipeline->cleanup(gfxDevice);

    if (shadowMapPipeline)
        shadowMapPipeline->cleanup(device);

    if (skyboxPipeline)
        skyboxPipeline->cleanup(device);

    if (meshPipeline)
        meshPipeline->cleanup(device);

    sceneDataBuffer.cleanup(gfxDevice);

    // if (drawImageId != NULL_IMAGE_ID)
    // gfxDevice.destroyImage();

    // if (depthImageId != NULL_IMAGE_ID)
    // gfxDevice.destroyImage(depthImageId);

    drawImageId = NULL_IMAGE_ID;
    depthImageId = NULL_IMAGE_ID;
    resources = nullptr;
    meshPipeline.reset();
    skyboxPipeline.reset();
    skinningPipeline.reset();
    lineRenderingPass.reset();
    renderer2DPass.reset();
    lightCullingPipeline.reset();
    shadowMapPipeline.reset();
    shadowLightDirection.reset();
    submittedLights.clear();
    spotShadows.clear();
    pendingMaterialUploads.clear();
    meshDrawCommands.clear();
    sortedMeshDrawCommands.clear();
    initialized = false;
    GameState::GetInstance().SetRenderer(nullptr);
}

void GameRenderer::drawMesh(MeshID id, const glm::mat4& transform, MaterialID materialId)
{
    const auto& mesh = resources->meshes().getMesh(id);
    const auto worldBoundingSphere = edge::calculateBoundingSphereWorld(transform, mesh.boundingSphere, false);
    assert(materialId != NULL_MATERIAL_ID);

    meshDrawCommands.push_back(MeshDrawCommand{
        .meshId = id,
        .transformMatrix = transform,
        .worldBoundingSphere = worldBoundingSphere,
        .materialId = materialId,
    });
}

void GameRenderer::drawSkinnedMesh(MeshID id,
                                   const glm::mat4& transform,
                                   MaterialID materialId,
                                   SkinnedMesh* skinnedMesh,
                                   std::size_t jointMatricesStartIndex)
{
    const auto& mesh = resources->meshes().getMesh(id);
    const auto worldBoundingSphere = edge::calculateBoundingSphereWorld(transform, mesh.boundingSphere, true);
    assert(materialId != NULL_MATERIAL_ID);
    assert(skinnedMesh != nullptr);

    meshDrawCommands.push_back(MeshDrawCommand{
        .meshId = id,
        .transformMatrix = transform,
        .worldBoundingSphere = worldBoundingSphere,
        .materialId = materialId,
        .skinnedMesh = skinnedMesh,
        .jointMatricesStartIndex = static_cast<std::uint32_t>(jointMatricesStartIndex),
    });
}

void GameRenderer::submitLight(const Light& light)
{
    submittedLights.push_back(light);
}

void GameRenderer::uploadLights(const Camera& camera)
{
    const Frustum cameraFrustum = edge::createFrustumFromCamera(camera);

    for (const Light& light : submittedLights)
    {
        lightCullingPipeline->appendLight(light, assignShadowIndex(light, cameraFrustum));
    }
}

std::uint32_t GameRenderer::assignShadowIndex(const Light& light, const Frustum& cameraFrustum)
{
    if (!light.castShadows)
    {
        return NO_SHADOW;
    }

    // there's a single directional shadow map, the first shadow casting directional light gets it
    if (light.type == LightType::Directional && !shadowLightDirection)
    {
        shadowLightDirection = light.direction;
        return 0;
    }

    if (light.type == LightType::Spot && spotShadows.size() < ShadowMapPipeline::MAX_SPOT_SHADOWS)
    {
        const GPUSpotShadow spotShadow = ShadowMapPipeline::calculateSpotShadow(light);
        if (!edge::frustumsOverlap(cameraFrustum, spotShadow.lightViewProj))
        {
            return NO_SHADOW;
        }

        spotShadows.push_back(spotShadow);
        return static_cast<std::uint32_t>(spotShadows.size() - 1);
    }

    return NO_SHADOW;
}

const GPUImage& GameRenderer::getDrawImage() const
{
    assert(resources);
    return resources->getImage(drawImageId);
}

Material& GameRenderer::getMaterialMutable(MaterialID id)
{
    assert(id != NULL_MATERIAL_ID);
    return resources->materials().getMaterialMutable(id);
}

void GameRenderer::updateMaterialGPU(MaterialID id)
{
    assert(id != NULL_MATERIAL_ID);
    pendingMaterialUploads.push_back(id);
}

void GameRenderer::setSkyboxTexture(ImageID skyboxImageId)
{
    spdlog::debug("Set skybox texture to image id {}", skyboxImageId);
    skyboxPipeline->setSkyboxImage(skyboxImageId);
}

void GameRenderer::recordMaterialUpdates(VkCommandBuffer cmd)
{
    const VkBuffer materialBuffer = resources->materials().getMaterialDataBuffer().buffer;

    if (!pendingMaterialUploads.empty())
    {
        // Updates go through the command buffer instead of the mapped pointer: the
        // copy then lands after earlier frames finished reading the old values,
        // without the CPU having to wait for them.
        vkutil::bufferShaderReadToTransferWriteBarrier(cmd, materialBuffer, 0, VK_WHOLE_SIZE);

        for (MaterialID id : pendingMaterialUploads)
        {
            const MaterialData data = resources->materials().getMaterialData(id);
            vkCmdUpdateBuffer(
                cmd,
                materialBuffer,
                static_cast<VkDeviceSize>(id) * sizeof(MaterialData),
                sizeof(MaterialData),
                &data);
        }
        pendingMaterialUploads.clear();
    }

    // newly added materials are written through the mapped pointer, updates by the copies above
    vkutil::bufferHostAndTransferWriteToShaderReadBarrier(cmd, materialBuffer, 0, VK_WHOLE_SIZE);
}

void GameRenderer::createDrawImage(GfxDevice& gfxDevice,
                                   const glm::ivec2& drawImageSize,
                                   bool firstCreate)
{
    const VkExtent3D drawImageExtent{
        .width = (std::uint32_t)drawImageSize.x,
        .height = (std::uint32_t)drawImageSize.y,
        .depth = 1,
    };

    constexpr VkSampleCountFlagBits noMsaa = VK_SAMPLE_COUNT_1_BIT;

    {
        // setup draw image (single-sampled)
        VkImageUsageFlags usages{};
        usages |= VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
        usages |= VK_IMAGE_USAGE_TRANSFER_DST_BIT;
        usages |= VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
        usages |= VK_IMAGE_USAGE_SAMPLED_BIT;

        auto createImageInfo = vkutil::CreateImageInfo{
            .format = drawImageFormat,
            .usage = usages,
            .extent = drawImageExtent,
            .samples = noMsaa,
        };

        // a fresh bindless slot, the old one may still be read by frames in flight
        resources->destroyImage(drawImageId);
        drawImageId = resources->createImage(gfxDevice, createImageInfo, "draw image");

        if (firstCreate)
        {
            // Optional: a separate post-fx target (ping-pong)
            // postFXDrawImageId = gfxDevice.createImage(createImageInfo, "post FX draw image");
        }
    }


    {
        // setup depth image (single-sampled)
        auto createInfo = vkutil::CreateImageInfo{
            .format = depthImageFormat,
            .usage = VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
            .extent = drawImageExtent,
            .samples = noMsaa,
        };

        resources->destroyImage(depthImageId);
        depthImageId = resources->createImage(gfxDevice, createInfo, "depth image");
        spdlog::info("Created depth image with id {}", depthImageId);
    }
}