//
// Created by atdunbg on 2026/9/23.
//

#include <Engine/Renderer/SceneRenderer.h>
#include <Engine/ECS/Scene.h>
#include <Engine/ECS/Components.h>
#include <Engine/Renderer/Camera.h>
#include <Engine/Renderer/Frustum.h>
#include <Engine/Renderer/Model.h>
#include <Engine/Renderer/Mesh.h>
#include <Engine/Assets/AssetManager.h>
#include <Engine/Core/Log.h>

#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>

#include <tracy/Tracy.hpp>

#include <chrono>
#include <limits>

namespace Engine
{
    using namespace ECS;   // TransformComponent / Entity 等直接使用

    namespace
    {
        // 相机裁剪范围（与 BuildFrame 中 GetProjection 的参数一致）
        constexpr f32 kCameraNear = 0.1f;
        constexpr f32 kCameraFar = 100.0f;
        // 实用分割插值因子（0=均匀，1=对数；0.6 是业界常用值）
        constexpr f32 kCascadeLambda = 0.6f;

        // Halton 低差异序列（TAA 子像素抖动）
        f32 Halton(u32 index, u32 base)
        {
            f32 f = 1.0f, r = 0.0f;
            while (index > 0) {
                f /= static_cast<f32>(base);
                r += f * static_cast<f32>(index % base);
                index /= base;
            }
            return r;
        }
    }

    void SceneRenderer::ComputeCascades(const glm::vec3& shadowDir)
    {
        constexpr u32 N = kMaxCascades;

        // ---- 1. 实用分割：均匀与对数分割按 λ 插值（视图深度，正值）----
        f32 splits[N];
        for (u32 i = 1; i <= N; ++i) {
            const f32 t = static_cast<f32>(i) / N;
            const f32 uniform = kCameraNear + (kCameraFar - kCameraNear) * t;
            const f32 log = kCameraNear * glm::pow(kCameraFar / kCameraNear, t);
            splits[i - 1] = glm::mix(uniform, log, kCascadeLambda);
        }
        for (u32 c = 0; c < N; ++c) {
            m_Frame.ubo.cascadeSplits[c] = splits[c];
        }

        // ---- 2. 全视锥 8 角反投影到世界（NDC z: 0=近 1=远，Vulkan [0,1]）----
        const glm::mat4 invVP = glm::inverse(m_Frame.ubo.projNoJitter * m_Frame.ubo.view);
        glm::vec3 nearC[4], farC[4];
        const glm::vec2 ndcXY[4] = {{-1,-1},{1,-1},{-1,1},{1,1}};
        for (u32 i = 0; i < 4; ++i) {
            glm::vec4 p = invVP * glm::vec4(ndcXY[i], 0.0f, 1.0f);
            nearC[i] = glm::vec3(p / p.w);
            p = invVP * glm::vec4(ndcXY[i], 1.0f, 1.0f);
            farC[i] = glm::vec3(p / p.w);
        }

        // ---- 3. 每级联：切片角点 → 光空间 AABB 拟合 → 纹素对齐 ----
        const glm::vec3 up = glm::abs(shadowDir.y) > 0.999f
            ? glm::vec3(0.0f, 0.0f, 1.0f) : glm::vec3(0.0f, 1.0f, 0.0f);

        for (u32 c = 0; c < N; ++c) {
            const f32 d0 = (c == 0) ? kCameraNear : splits[c - 1];
            const f32 d1 = splits[c];
            const f32 t0 = (d0 - kCameraNear) / (kCameraFar - kCameraNear);
            const f32 t1 = (d1 - kCameraNear) / (kCameraFar - kCameraNear);

            glm::vec3 slice[8];
            for (u32 i = 0; i < 4; ++i) {
                slice[i]     = glm::mix(nearC[i], farC[i], t0);
                slice[4 + i] = glm::mix(nearC[i], farC[i], t1);
            }

            // 光视图：位于切片中心，沿光方向看
            glm::vec3 center(0.0f);
            for (const auto& p : slice) center += p;
            center *= (1.0f / 8.0f);
            const glm::mat4 lightView = glm::lookAt(center, center + shadowDir, up);

            glm::vec3 mn( std::numeric_limits<f32>::max());
            glm::vec3 mx( std::numeric_limits<f32>::lowest());
            for (const auto& p : slice) {
                const glm::vec3 lv = glm::vec3(lightView * glm::vec4(p, 1.0f));
                mn = glm::min(mn, lv);
                mx = glm::max(mx, lv);
            }

            // 沿光方向扩展深度范围：切片盒外的 caster（如高树、城墙）
            // 也能投影子进本级联——per-cascade 剔除的正确性前提。
            // 只扩 z 不缩 XY，阴影分辨率不受影响
            constexpr f32 kCasterPad = 30.0f;
            mn.z -= kCasterPad;
            mx.z += kCasterPad;

            // 纹素对齐：X/Y 边界量化到纹素网格（消除相机移动时的阴影边缘闪烁）
            const f32 texelX = (mx.x - mn.x) / static_cast<f32>(kCascadeRes);
            const f32 texelY = (mx.y - mn.y) / static_cast<f32>(kCascadeRes);
            mn.x = glm::floor(mn.x / texelX) * texelX;
            mn.y = glm::floor(mn.y / texelY) * texelY;
            mx.x = mn.x + glm::ceil((mx.x - mn.x) / texelX) * texelX;
            mx.y = mn.y + glm::ceil((mx.y - mn.y) / texelY) * texelY;

            // 光视图空间里 z 为负（看向 -Z）；近远距离取正
            glm::mat4 lightProj = glm::ortho(mn.x, mx.x, mn.y, mx.y,
                                             -mx.z, -mn.z);
            // Vulkan clip 深度 [0,1]（ortho 的 w_clip=1，偏移折进 z 平移分量）
            lightProj[2][2] *= 0.5f;
            lightProj[3][2] = lightProj[3][2] * 0.5f + 0.5f;

            m_Frame.ubo.lightSpace[c] = lightProj * lightView;
        }
    }

    bool SceneRenderer::BuildFrame(ECS::Scene& scene, Assets::AssetManager& assets,
                                   vk::Extent2D targetExtent)
    {
        ZoneScopedN("BuildFrame");
        const auto collectStart = std::chrono::steady_clock::now();

        // ---- 0. TAA 时域状态先行（velocity = 上一帧姿态与相机）----
        UniformBufferObject& ubo = m_Frame.ubo;
        ubo.prevView = m_LastView;
        ubo.prevProj = m_LastProjNoJitter;

        // ---- 1. 相机（位姿单一事实源 = 相机实体的 Transform）----
        TransformComponent* camTransform = nullptr;
        auto* cameraComp = scene.GetPrimaryCamera(&camTransform);
        if (!cameraComp || !camTransform) {
            // 无相机：本帧跳过（正常场景必有相机，安全网）
            return false;
        }
        const Camera& camera = cameraComp->camera;

        const f32 aspect = static_cast<f32>(targetExtent.width) /
                           static_cast<f32>(targetExtent.height);
        ubo.view = camera.GetView(*camTransform);
        ubo.proj = camera.GetProjection(aspect, 0.1f, 100.0f);

        // 未抖动投影：velocity 的当前帧 clip 位置用它（抖动不能进 velocity）
        ubo.projNoJitter = ubo.proj;

        // 子像素抖动（Halton×8）：光栅化采样位置逐帧偏移，TAA 累积去锯齿。
        // NDC 偏移 = 像素偏移 × 2 / 尺寸
        {
            const u32 idx = static_cast<u32>(m_FrameIndex % 8);
            const f32 jx = Halton(idx + 1, 2) - 0.5f;
            const f32 jy = Halton(idx + 1, 3) - 0.5f;
            ubo.proj[2][0] += jx * 2.0f / static_cast<f32>(targetExtent.width);
            ubo.proj[2][1] += jy * 2.0f / static_cast<f32>(targetExtent.height);
        }

        ubo.ambientColor = scene.GetSettings().ambientColor;

        // 天空盒：视图去平移（只留旋转）× 投影 → 逆矩阵，push 给 Skybox pass；
        // 反投影远平面点 = 相机相对世界方向，frag normalize 后采样 cube
        m_Frame.skyInvVP =
            glm::inverse(ubo.projNoJitter * glm::mat4(glm::mat3(ubo.view)));

        // ---- 2. 灯光（第一个方向光作为 CSM 阴影光源；点光位置来自 Transform）----
        u32 lightCount = 0;
        glm::vec3 shadowDir{0.0f, -1.0f, 0.0f};   // 无灯兜底：头顶正射
        scene.Each<TransformComponent, LightComponent>(
            [&](Entity, TransformComponent& t, LightComponent& light) {
            if (lightCount >= kMaxLights) return;

            auto& dst = ubo.lights[lightCount];
            dst.color = light.color;
            dst.params = glm::vec4(
                light.type == LightComponent::Type::Point ? 1.0f : 0.0f,
                light.radius, 0.0f, 0.0f);

            if (light.type == LightComponent::Type::Point) {
                dst.position = glm::vec4(t.GetPosition(), 1.0f);
                dst.direction = glm::vec4(0.0f);
            } else {
                glm::vec3 dir = glm::length(light.direction) > 0.0f
                    ? glm::normalize(light.direction) : glm::vec3(0.0f);
                dst.direction = glm::vec4(dir, 0.0f);
                dst.position = glm::vec4(0.0f);
            }

            // CSM 阴影光源 = 第一个方向光（点光不做级联阴影）
            if (lightCount == 0 &&
                light.type == LightComponent::Type::Directional) {
                shadowDir = glm::length(light.direction) > 0.0f
                    ? glm::normalize(light.direction) : shadowDir;
            }
            ++lightCount;
        });
        for (u32 i = lightCount; i < kMaxLights; ++i) {
            ubo.lights[i] = {};
        }
        ComputeCascades(shadowDir);
        ubo.cameraPos = glm::vec4(camTransform->GetPosition(), 1.0f);

        // ---- 3. GPUDriven：单趟收集 + 多视图剔除 ----
        // 同作用域第二个 zone：ZoneScopedN 变量名固定会重声明，用 ZoneNamedN
        ZoneNamedN( ___tracy_collect, "GPUDriven Collect", true );
        // 每个实例的包围球对 4 个视锥分别测试：
        //   相机视锥 → 主视图（Forward 画）
        //   各级联光视锥 → 各自的阴影 batch（Shadow 按 region 各画各的）
        // 级联盒沿光方向加了 caster pad（见 ComputeCascades）：
        // 盒外 caster 投进级联的影子不会被剔除裁掉。
        // 收集按视图分桶，最后拼接成 [主段|级联0|级联1|级联2]，
        // batch.first 拼接时加上段基址，保证段内区间连续。
        m_Frame.instanceModels.clear();
        m_Frame.prevInstanceModels.clear();
        m_MainScratch.Clear();
        for (auto& s : m_CascScratch) s.Clear();
        u32 totalInstances = 0;
        bool instanceTruncated = false;

        // 剔除平面提取按 [0,1] 深度约定——相机 proj 与 lightSpace 都已
        // 做 [0,1] 重映射；原始 [-1,1] 矩阵的 near 平面公式不同，勿用
        const Frustum mainFrustum = Frustum::FromMatrix(ubo.proj * ubo.view);
        Frustum cascFrustum[kMaxCascades];
        for (u32 c = 0; c < kMaxCascades; ++c) {
            cascFrustum[c] = Frustum::FromMatrix(ubo.lightSpace[c]);
        }

        scene.Each<TransformComponent, MeshRendererComponent>(
            [&](Entity entity, TransformComponent& t, MeshRendererComponent& mr) {
                Model* model = assets.GetModel(mr.model);
                if (!model) return;

                // TAA velocity：本实例的上一帧世界矩阵（首帧/新实体 = 当前
                // 矩阵 → 速度为 0，不产生假运动）
                const u32 entKey = static_cast<u32>(entt::to_integral(entity));
                glm::mat4 prevWorld = t.World();
                if (auto it = m_PrevWorlds.find(entKey); it != m_PrevWorlds.end()) {
                    prevWorld = it->second;
                }
                m_CurWorlds.emplace(entKey, t.World());

                // 实例世界缩放（矩阵列向量最大长度，包围球半径随缩放）
                const f32 worldScale = glm::max(
                    glm::length(glm::vec3(t.World()[0])),
                    glm::max(glm::length(glm::vec3(t.World()[1])),
                             glm::length(glm::vec3(t.World()[2]))));

                for (auto& mesh : model->GetMeshes()) {
                    const auto& bounds = mesh->GetBounds();
                    const glm::vec3 worldCenter =
                        glm::vec3(t.World() * glm::vec4(bounds.center, 1.0f));
                    const f32 worldRadius = bounds.radius * worldScale;

                    auto tryAppend = [&](ViewScratch& s, const Frustum& frustum) {
                        if (!frustum.IntersectsSphere(worldCenter, worldRadius)) {
                            return;
                        }
                        if (totalInstances >= kMaxInstances) {
                            instanceTruncated = true;
                            return;
                        }
                        const auto [it, inserted] = s.lookup.emplace(
                            mesh.get(), static_cast<u32>(s.worlds.size()));
                        if (inserted) {
                            s.result.batches.push_back({mesh.get(),
                                static_cast<u32>(s.worlds.size()), 0});
                        }
                        s.result.batches[it->second].count++;
                        s.result.instanceCount++;
                        s.worlds.push_back(t.World());
                        s.prevWorlds.push_back(prevWorld);
                        ++totalInstances;
                    };

                    tryAppend(m_MainScratch, mainFrustum);
                    for (u32 c = 0; c < kMaxCascades; ++c) {
                        tryAppend(m_CascScratch[c], cascFrustum[c]);
                    }
                }
            });

        // ---- 4. 拼接：[主段|级联0|级联1|级联2]，batch.first 加段基址 ----
        auto appendView = [&](CullResult& dst, ViewScratch& s) {
            const u32 base = static_cast<u32>(m_Frame.instanceModels.size());
            dst.batches.clear();
            dst.batches.reserve(s.result.batches.size());
            for (auto& b : s.result.batches) {
                b.first += base;
                dst.batches.push_back(b);
            }
            dst.instanceCount = s.result.instanceCount;
            m_Frame.instanceModels.insert(m_Frame.instanceModels.end(),
                                          s.worlds.begin(), s.worlds.end());
            m_Frame.prevInstanceModels.insert(m_Frame.prevInstanceModels.end(),
                                              s.prevWorlds.begin(),
                                              s.prevWorlds.end());
        };
        appendView(m_Frame.mainView, m_MainScratch);
        for (u32 c = 0; c < kMaxCascades; ++c) {
            appendView(m_Frame.shadowCascades[c], m_CascScratch[c]);
        }

        // 帧末滚动时域状态：cur → prev（每帧重建，销毁的实体自动清出）
        m_PrevWorlds.swap(m_CurWorlds);
        m_CurWorlds.clear();
        m_LastView = ubo.view;
        m_LastProjNoJitter = ubo.projNoJitter;
        ++m_FrameIndex;

        if (instanceTruncated && !m_TruncationWarned) {
            // 只告警一次（每帧告警会刷屏）
            m_TruncationWarned = true;
            ENGINE_WARN("GPUDriven: instance buffer full ({}), entities truncated",
                        kMaxInstances);
        }

        m_Frame.collectMs = std::chrono::duration<f32, std::milli>(
            std::chrono::steady_clock::now() - collectStart).count();
        return true;
    }

} // namespace Engine