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

#include <Engine/Renderer/Passes/ShadowPass.h>
#include <Engine/Renderer/Vertex.h>
#include <Engine/Renderer/Mesh.h>

#include <array>

namespace Engine::Passes
{
    void ShadowPass::Register(RG::RenderGraph& rg)
    {
        // 所有渲染体都写入每个级联，不做实体级开关——castShadow 应由
        // 材质/实体数据驱动剔除，而非 pass 逻辑特判。
        // （每级联独立剔除是后续优化点：shadowView 细分为 per-cascade 列表）
        rg.AddPass("Shadow", [this](RG::RenderGraphContext& ctx) {
            const vk::CommandBuffer cmd = ctx.cmd;

            for (u32 c = 0; c < kMaxCascades; ++c) {
                vk::RenderingAttachmentInfo depthAttachment;
                depthAttachment.setImageView(ctx.View(m_ShadowMaps[c]))
                               .setImageLayout(vk::ImageLayout::eDepthAttachmentOptimal)
                               .setLoadOp(vk::AttachmentLoadOp::eClear)
                               .setStoreOp(vk::AttachmentStoreOp::eStore)
                               .setClearValue(vk::ClearValue(
                                   vk::ClearDepthStencilValue(1.0f, 0)));

                vk::RenderingInfo renderingInfo;
                renderingInfo.setRenderArea({{0, 0}, {kCascadeRes, kCascadeRes}})
                             .setLayerCount(1)
                             .setPDepthAttachment(&depthAttachment);

                cmd.beginRendering(renderingInfo);

                std::array<vk::ShaderStageFlagBits, 2> stages = {
                    vk::ShaderStageFlagBits::eVertex,
                    vk::ShaderStageFlagBits::eFragment
                };
                std::array<vk::ShaderEXT, 2> shaders = {m_Gpu.shadowVert, m_Gpu.shadowFrag};
                cmd.bindShadersEXT(stages, shaders);

                // ---- 动态状态（depth-only）----
                cmd.setRasterizerDiscardEnable(VK_FALSE);

                cmd.setDepthTestEnable(VK_TRUE);
                cmd.setDepthWriteEnable(VK_TRUE);
                cmd.setDepthCompareOp(vk::CompareOp::eLess);
                cmd.setDepthBoundsTestEnable(VK_FALSE);

                cmd.setStencilTestEnable(VK_FALSE);
                // 深度 bias 压平 shadow acne（NDC 量级）
                cmd.setDepthBiasEnable(VK_TRUE);
                cmd.setDepthBias(0.0025f, 0.0f, 0.003f);
                cmd.setPolygonModeEXT(vk::PolygonMode::eFill);
                cmd.setCullMode(vk::CullModeFlagBits::eNone);
                cmd.setFrontFace(vk::FrontFace::eCounterClockwise);
                cmd.setPrimitiveTopology(vk::PrimitiveTopology::eTriangleList);
                cmd.setPrimitiveRestartEnable(VK_FALSE);

                auto binding = Vertex::GetBindingDescription();
                auto attributes = Vertex::GetAttributeDescriptions();
                cmd.setVertexInputEXT(binding, attributes);

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

                vk::Viewport viewport(0.0f, 0.0f,
                                      static_cast<float>(kCascadeRes),
                                      static_cast<float>(kCascadeRes),
                                      0.0f, 1.0f);
                cmd.setViewportWithCount(viewport);

                vk::Rect2D scissor({0, 0}, {kCascadeRes, kCascadeRes});
                cmd.setScissorWithCount(scissor);

                // ---- 绘制：GPUDriven——与 Forward 共用实例数据，但用
                //      光源视锥剔除后的 batch（相机外实体仍投影子）----
                cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics,
                                       m_Gpu.shadowLayout, 0,
                                       m_Gpu.uboSets[ctx.frame.frameSlot], {});

                // stage 必须覆盖 layout range 声明的全部 stage（VUID-01796）
                // 每级联 push 各自的光空间 VP
                cmd.pushConstants(m_Gpu.shadowLayout,
                                  vk::ShaderStageFlagBits::eVertex |
                                  vk::ShaderStageFlagBits::eFragment,
                                  0, sizeof(glm::mat4),
                                  &m_Frame.ubo.lightSpace[c]);

                for (const auto& batch : m_Frame.shadowCascades[c].batches) {
                    batch.mesh->Draw(cmd, batch.count, batch.first);
                }

                cmd.endRendering();
            }
        })
        .Write(m_ShadowMaps[0], RG::ResourceUse::DepthAttachment)
        .Write(m_ShadowMaps[1], RG::ResourceUse::DepthAttachment)
        .Write(m_ShadowMaps[2], RG::ResourceUse::DepthAttachment);
    }

    void ShadowPass::OnTargetsRebuilt(RG::RenderGraph& rg)
    {
        m_Gpu.BindShadowMaps(rg, m_ShadowMaps);
    }

} // namespace Engine::Passes