#include <destrum/Graphics/Renderer.h>

#include <destrum/Graphics/Managers/LineRenderingManager.h>
#include <destrum/Graphics/Managers/QuadRenderingManager.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);

        quadRendererPass = std::make_unique<QuadRendererPass>();
        quadRendererPass->init(gfxDevice, _resources, drawImageFormat);

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

void GameRenderer::beginDrawing(GfxDevice& gfxDevice)
{
    flushMaterialUpdates(gfxDevice);
    meshDrawCommands.clear();
    skinningPipeline->beginDrawing(gfxDevice.getCurrentFrameIndex());
    lightCullingPipeline->beginDrawing(gfxDevice.getCurrentFrameIndex());
    shadowLightDirection.reset();
    QuadRenderingManager::GetInstance().ClearQuads();
}

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);

        lightCullingPipeline->flushCurrentFrame(cmd, gfxDevice);
    }

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

    // 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,
    };

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

        vkutil::bufferHostWriteToShaderReadBarrier(
            cmd,
            resources->materials().getMaterialDataBuffer().buffer,
            0,
            VK_WHOLE_SIZE
        );
    }

    {
        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->draw(
            cmd,
            *resources,
            sceneDataBuffer.getBuffer(),
            meshDrawCommands,
            sortedMeshDrawCommands);
    }

    {
        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);
    }

    {
        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, "QuadRendererPass::draw");
        vkutil::ScopedCmdLabel label(cmd, "QuadRendererPass::draw", vkutil::labelColors::ui);

        quadRendererPass->draw(
            cmd,
            gfxDevice,
            *resources,
            drawImage,
            glm::ortho(0.0f,
                       static_cast<float>(drawImage.extent.width),
                       static_cast<float>(drawImage.extent.height),
                       0.0f, -1.0f, 1.0f));
    }
}

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

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

    if (device != VK_NULL_HANDLE)
    {
        vkDeviceWaitIdle(device);
    }

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

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

    if (quadRendererPass)
        quadRendererPass->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();
    quadRendererPass.reset();
    lightCullingPipeline.reset();
    shadowMapPipeline.reset();
    shadowLightDirection.reset();
    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)
{
    Light submittedLight = light;

    // there's a single shadow map, the first shadow casting directional light gets it
    submittedLight.castShadows =
        light.castShadows && light.type == LightType::Directional && !shadowLightDirection;

    if (submittedLight.castShadows)
    {
        shadowLightDirection = light.direction;
    }

    lightCullingPipeline->appendLight(submittedLight);
}

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::flushMaterialUpdates(GfxDevice& gfxDevice)
{
    if (!pendingMaterialUploads.empty())
    {
        // Material data is currently shared by all frame contexts. Wait before
        // mutating it so an older in-flight frame cannot read the overwritten
        // range. This can later be replaced with per-frame material buffers.
        gfxDevice.waitIdle();
    }

    for (MaterialID id : pendingMaterialUploads)
    {
        resources->materials().updateMaterialGPU(id);
        gfxDevice.getMemoryManager().flushAllocation(
            resources->materials().getMaterialDataBuffer(),
            static_cast<VkDeviceSize>(id) * sizeof(MaterialData),
            sizeof(MaterialData));
    }
    pendingMaterialUploads.clear();
}

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,
        };

        // reuse the same id if creating again
        drawImageId = resources->createImage(gfxDevice, createImageInfo, "draw image", nullptr, drawImageId);

        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,
        };

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