//
// EditorGizmoPass.cpp：pass 装配（shader object）与录制。
// editor 专属（HELLO_EDITOR）：game 构建整 TU 剔除
//（Engine/CMakeLists REMOVE_ITEM）。
//

#include <Engine/Renderer/Passes/EditorGizmoPass.h>
#include <Engine/Vulkan/Device.h>
#include <Engine/Vulkan/DescriptorManager.h>
#include <Engine/Vulkan/UniformBuffer.h>
#include <Engine/Core/Log.h>

#include <glm/glm.hpp>

#include <tracy/Tracy.hpp>

#include <algorithm>
#include <array>
#include <cstddef>
#include <string>

namespace Engine::Passes
{
    namespace
    {
        // 线流布局：与 debug_line 同（location0 pos / location1 color）。
        // setVertexInputEXT 走 2EXT 结构（0 divisor 需 feature，不用）
        vk::VertexInputBindingDescription2EXT MakeWireBinding()
        {
            vk::VertexInputBindingDescription2EXT b;
            b.setBinding(0)
             .setStride(sizeof(GizmoWireVertex))
             .setInputRate(vk::VertexInputRate::eVertex)
             .setDivisor(1);
            return b;
        }

        std::array<vk::VertexInputAttributeDescription2EXT, 2>
        MakeWireAttributes()
        {
            std::array<vk::VertexInputAttributeDescription2EXT, 2> a;
            a[0].setLocation(0).setBinding(0)
                .setFormat(vk::Format::eR32G32B32Sfloat)
                .setOffset(offsetof(GizmoWireVertex, pos));
            a[1].setLocation(1).setBinding(0)
                .setFormat(vk::Format::eR32G32B32A32Sfloat)
                .setOffset(offsetof(GizmoWireVertex, color));
            return a;
        }

        // 图标流布局：pos / color / aux（shapeId + uv）。entityId 槽位
        // 本 pass 不消费——拾取由 HitProxyPass 用同一 VB 全四属性画
        vk::VertexInputBindingDescription2EXT MakeIconBinding()
        {
            vk::VertexInputBindingDescription2EXT b;
            b.setBinding(0)
             .setStride(sizeof(GizmoQuadVertex))
             .setInputRate(vk::VertexInputRate::eVertex)
             .setDivisor(1);
            return b;
        }

        std::array<vk::VertexInputAttributeDescription2EXT, 3>
        MakeIconAttributes()
        {
            std::array<vk::VertexInputAttributeDescription2EXT, 3> a;
            a[0].setLocation(0).setBinding(0)
                .setFormat(vk::Format::eR32G32B32Sfloat)
                .setOffset(offsetof(GizmoQuadVertex, pos));
            a[1].setLocation(1).setBinding(0)
                .setFormat(vk::Format::eR32G32B32A32Sfloat)
                .setOffset(offsetof(GizmoQuadVertex, color));
            a[2].setLocation(2).setBinding(0)
                .setFormat(vk::Format::eR32G32B32A32Sfloat)
                .setOffset(offsetof(GizmoQuadVertex, aux));
            return a;
        }

        // 直写方程（线流）：(ONE, ZERO) → dst = src，数学上等价无混合
        vk::ColorBlendEquationEXT MakeDirectEquation()
        {
            vk::ColorBlendEquationEXT eq;
            eq.setSrcColorBlendFactor(vk::BlendFactor::eOne)
              .setDstColorBlendFactor(vk::BlendFactor::eZero)
              .setColorBlendOp(vk::BlendOp::eAdd)
              .setSrcAlphaBlendFactor(vk::BlendFactor::eOne)
              .setDstAlphaBlendFactor(vk::BlendFactor::eZero)
              .setAlphaBlendOp(vk::BlendOp::eAdd);
            return eq;
        }

        // 预乘 alpha（图标流）：shader 已输出预乘色，src 原样加权叠加
        vk::ColorBlendEquationEXT MakePremultEquation()
        {
            vk::ColorBlendEquationEXT eq;
            eq.setSrcColorBlendFactor(vk::BlendFactor::eOne)
              .setDstColorBlendFactor(vk::BlendFactor::eOneMinusSrcAlpha)
              .setColorBlendOp(vk::BlendOp::eAdd)
              .setSrcAlphaBlendFactor(vk::BlendFactor::eOne)
              .setDstAlphaBlendFactor(vk::BlendFactor::eOneMinusSrcAlpha)
              .setAlphaBlendOp(vk::BlendOp::eAdd);
            return eq;
        }
    } // namespace

    EditorGizmoPass::EditorGizmoPass(Device& device, ShaderManager& shaders,
                                     DescriptorManager& descriptors,
                                     FrameSceneData& frame,
                                     RG::TextureHandle hdrColor,
                                     RG::TextureHandle depth,
                                     UniformBuffer& iconVB,
                                     UniformBuffer& wireVB)
        : m_Device(device), m_Descriptors(descriptors), m_Frame(frame),
          m_HDRColor(hdrColor), m_Depth(depth), m_IconVB(iconVB),
          m_WireVB(wireVB)
    {
        // 无描述符 set：viewProj 走 push constant（两对 shader 同契约）
        m_Layout = descriptors.CreatePipelineLayout(
            std::vector<vk::DescriptorSetLayout>{}, sizeof(glm::mat4));

        const std::string shaderDir = "shaders/";
        auto wireVert = shaders.CreateFromFile(
            shaderDir + "debug_line.vert.spv",
            vk::ShaderStageFlagBits::eVertex,
            std::vector<vk::DescriptorSetLayout>{},
            sizeof(glm::mat4));
        auto wireFrag = shaders.CreateFromFile(
            shaderDir + "debug_line.frag.spv",
            vk::ShaderStageFlagBits::eFragment,
            std::vector<vk::DescriptorSetLayout>{},
            sizeof(glm::mat4));
        auto iconVert = shaders.CreateFromFile(
            shaderDir + "gizmo_billboard.vert.spv",
            vk::ShaderStageFlagBits::eVertex,
            std::vector<vk::DescriptorSetLayout>{},
            sizeof(glm::mat4));
        auto iconFrag = shaders.CreateFromFile(
            shaderDir + "gizmo_billboard.frag.spv",
            vk::ShaderStageFlagBits::eFragment,
            std::vector<vk::DescriptorSetLayout>{},
            sizeof(glm::mat4));

        m_WireVert = std::move(wireVert);
        m_WireFrag = std::move(wireFrag);
        m_IconVert = std::move(iconVert);
        m_IconFrag = std::move(iconFrag);
        if (!m_WireVert || !m_WireFrag || !m_IconVert || !m_IconFrag) {
            ENGINE_FATAL("EditorGizmoPass: failed to create shader objects");
        }
    }

    EditorGizmoPass::~EditorGizmoPass()
    {
        // shader 随 ShaderRef 引用计数自动销毁，此处只清其余原生句柄
        m_Device.GetDevice().destroyPipelineLayout(m_Layout);
    }

    void EditorGizmoPass::Register(RG::RenderGraph& rg)
    {
        rg.AddPass("EditorGizmo", [this](RG::RenderGraphContext& ctx) {
            const u32 iconCount = static_cast<u32>(
                m_Frame.gizmoIconVerts.size());
            const u32 wireCount = static_cast<u32>(
                m_Frame.gizmoWireVerts.size());
            // 骨骼线（选中实体骨架）：与 gizmo 线同 VB 同管线，画在
            // gizmo 段之后。截断公式与 Renderer 上传端一致（同一 VB
            // 容量推导，天然对齐）
            const u64 wireVBVerts =
                m_WireVB.GetSize() / sizeof(GizmoWireVertex);
            const u64 boneRoom =
                wireVBVerts > wireCount ? wireVBVerts - wireCount : 0;
            const u32 boneCount = static_cast<u32>(std::min<u64>(
                m_Frame.boneWireVerts.size(), boneRoom));
            if (iconCount == 0 && wireCount == 0 && boneCount == 0) {
                return;   // 空帧零开销
            }

            ZoneScopedN("EditorGizmo");
            const vk::CommandBuffer cmd = ctx.cmd;

            // HDR：Load 承接 Forward/PhysicsDebug 结果叠加；深度：Load
            // 承接场景深度做测试（只测不写），DontCare——HitProxy 稍后
            // Clear 重建
            vk::RenderingAttachmentInfo color;
            color.setImageView(ctx.View(m_HDRColor))
                 .setImageLayout(vk::ImageLayout::eColorAttachmentOptimal)
                 .setLoadOp(vk::AttachmentLoadOp::eLoad)
                 .setStoreOp(vk::AttachmentStoreOp::eStore);

            vk::RenderingAttachmentInfo depth;
            depth.setImageView(ctx.View(m_Depth))
                 .setImageLayout(vk::ImageLayout::eDepthAttachmentOptimal)
                 .setLoadOp(vk::AttachmentLoadOp::eLoad)
                 .setStoreOp(vk::AttachmentStoreOp::eDontCare);

            vk::RenderingInfo renderingInfo;
            renderingInfo.setRenderArea({{0, 0}, ctx.extent})
                         .setLayerCount(1)
                         .setColorAttachments(color)
                         .setPDepthAttachment(&depth);

            cmd.beginRendering(renderingInfo);

            const std::array<vk::ShaderStageFlagBits, 2> stages = {
                vk::ShaderStageFlagBits::eVertex,
                vk::ShaderStageFlagBits::eFragment
            };

            // layout 的 push range 是 vertex|fragment（CreatePipelineLayout
            // 默认），stageFlags 必须覆盖 range 全部 stage（VUID-01796）
            cmd.pushConstants<glm::mat4>(
                m_Layout,
                vk::ShaderStageFlagBits::eVertex |
                    vk::ShaderStageFlagBits::eFragment,
                0, m_Frame.ubo.proj * m_Frame.ubo.view);

            // ---- 公共动态状态（shader object 无管线对象，无条件直写；
            //      跨 pass 残留教训：公共状态块不进条件分支）----
            cmd.setRasterizerDiscardEnable(VK_FALSE);

            cmd.setDepthTestEnable(VK_TRUE);
            cmd.setDepthCompareOp(vk::CompareOp::eLessOrEqual);
            cmd.setDepthWriteEnable(VK_FALSE);   // gizmo 不污染场景深度
            cmd.setDepthBoundsTestEnable(VK_FALSE);
            cmd.setStencilTestEnable(VK_FALSE);
            cmd.setDepthBiasEnable(VK_FALSE);

            cmd.setPolygonModeEXT(vk::PolygonMode::eFill);
            cmd.setCullMode(vk::CullModeFlagBits::eNone);
            cmd.setFrontFace(vk::FrontFace::eCounterClockwise);
            cmd.setPrimitiveRestartEnable(VK_FALSE);

            cmd.setRasterizationSamplesEXT(vk::SampleCountFlagBits::e1);
            constexpr vk::SampleMask sampleMask = 0xFFFFFFFF;
            cmd.setSampleMaskEXT(vk::SampleCountFlagBits::e1, sampleMask);
            cmd.setAlphaToCoverageEnableEXT(VK_FALSE);
            cmd.setAlphaToOneEnableEXT(VK_FALSE);

            const vk::Viewport viewport(
                0.0f, 0.0f, static_cast<float>(ctx.extent.width),
                static_cast<float>(ctx.extent.height), 0.0f, 1.0f);
            cmd.setViewportWithCount(viewport);
            cmd.setScissorWithCount(vk::Rect2D({0, 0}, ctx.extent));

            constexpr vk::ColorComponentFlags colorWriteMask =
                vk::ColorComponentFlagBits::eR |
                vk::ColorComponentFlagBits::eG |
                vk::ColorComponentFlagBits::eB |
                vk::ColorComponentFlagBits::eA;
            cmd.setColorWriteMaskEXT(0, colorWriteMask);

            // ---- 线流：视锥线框 + 选中实体骨骼（LineList 无混合直写，
            //      同 VB 两段：[gizmo 线 | 骨骼线]）----
            if (wireCount > 0 || boneCount > 0) {
                const std::array<vk::ShaderEXT, 2> wireShaders = {
                    m_WireVert.GetHandle(), m_WireFrag.GetHandle()};
                cmd.bindShadersEXT(stages, wireShaders);

                const vk::VertexInputBindingDescription2EXT binding =
                    MakeWireBinding();
                const auto attrs = MakeWireAttributes();
                cmd.setVertexInputEXT(binding, attrs);

                cmd.setPrimitiveTopology(
                    vk::PrimitiveTopology::eLineList);
                cmd.setLineWidth(1.0f);   // 线拓扑必须显式（VUID-08617）

                constexpr vk::Bool32 blendOff = VK_FALSE;
                const vk::ColorBlendEquationEXT direct =
                    MakeDirectEquation();
                cmd.setColorBlendEnableEXT(0, blendOff);
                cmd.setColorBlendEquationEXT(0, direct);

                const std::array<vk::Buffer, 1> bufs{
                    m_WireVB.Get(ctx.frame.frameSlot)};
                const std::array<vk::DeviceSize, 1> offs{0};
                cmd.bindVertexBuffers(0, bufs, offs);
                if (wireCount > 0) {
                    cmd.draw(wireCount, 1, 0, 0);
                }
                if (boneCount > 0) {
                    cmd.draw(boneCount, 1, wireCount, 0);
                }
            }

            // ---- 图标流：billboard（TriangleList 预乘 alpha 叠加）----
            if (iconCount > 0) {
                const std::array<vk::ShaderEXT, 2> iconShaders = {
                    m_IconVert.GetHandle(), m_IconFrag.GetHandle()};
                cmd.bindShadersEXT(stages, iconShaders);

                const vk::VertexInputBindingDescription2EXT binding =
                    MakeIconBinding();
                const auto attrs = MakeIconAttributes();
                cmd.setVertexInputEXT(binding, attrs);

                cmd.setPrimitiveTopology(
                    vk::PrimitiveTopology::eTriangleList);

                constexpr vk::Bool32 blendOn = VK_TRUE;
                const vk::ColorBlendEquationEXT premult =
                    MakePremultEquation();
                cmd.setColorBlendEnableEXT(0, blendOn);
                cmd.setColorBlendEquationEXT(0, premult);

                const std::array<vk::Buffer, 1> bufs{
                    m_IconVB.Get(ctx.frame.frameSlot)};
                const std::array<vk::DeviceSize, 1> offs{0};
                cmd.bindVertexBuffers(0, bufs, offs);
                cmd.draw(iconCount, 1, 0, 0);
            }

            cmd.endRendering();
        })
        // Write use（WAW 链）：HDR 排 Forward/PhysicsDebug 之后、TAA 之前；
        // Depth WAW 把本 pass 固定排进深度链（HitProxy 清深度之前——
        // gizmo 深度测试消费场景深度，不是上帧残留）
        .Write(m_HDRColor, RG::ResourceUse::ColorAttachment)
        .Write(m_Depth, RG::ResourceUse::DepthAttachment);
    }

} // namespace Engine::Passes