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

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

#include <array>
#include <span>

namespace Engine::Passes
{
    void ForwardPass::Register(RG::RenderGraph& rg)
    {
        rg.AddPass("Forward", [this](RG::RenderGraphContext& ctx) {
            const vk::CommandBuffer cmd = ctx.cmd;

            vk::ClearValue clearValue;
            clearValue.color = vk::ClearColorValue(
                std::array<float, 4>{0.1f, 0.1f, 0.1f, 0.0f});

            vk::RenderingAttachmentInfo colorAttachments[2];
            // HDR 附件由 Skybox pass 清屏并画背景，这里 Load 在其上合成几何
            colorAttachments[0].setImageView(ctx.View(m_HDRColor))
                               .setImageLayout(vk::ImageLayout::eColorAttachmentOptimal)
                               .setLoadOp(vk::AttachmentLoadOp::eLoad)
                               .setStoreOp(vk::AttachmentStoreOp::eStore)
                               .setClearValue(clearValue);

            // velocity 附件：清 0 = 静止（背景无几何 → 历史原位采样）
            vk::ClearValue velocityClear;
            velocityClear.color = vk::ClearColorValue(std::array<float, 4>{0.f, 0.f, 0.f, 0.f});
            colorAttachments[1].setImageView(ctx.View(m_Velocity))
                               .setImageLayout(vk::ImageLayout::eColorAttachmentOptimal)
                               .setLoadOp(vk::AttachmentLoadOp::eClear)
                               .setStoreOp(vk::AttachmentStoreOp::eStore)
                               .setClearValue(velocityClear);

            vk::RenderingAttachmentInfo depthAttachment;
            depthAttachment.setImageView(ctx.View(m_Depth))
                           .setImageLayout(vk::ImageLayout::eDepthAttachmentOptimal)
                           .setLoadOp(vk::AttachmentLoadOp::eClear)
                           .setStoreOp(vk::AttachmentStoreOp::eDontCare)
                           .setClearValue(vk::ClearValue(
                               vk::ClearDepthStencilValue(1.0f, 0)));

            std::array<vk::RenderingAttachmentInfo, 2> colorArray = {
                colorAttachments[0], colorAttachments[1]};

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

            cmd.beginRendering(renderingInfo);

            // 绑定 shader object
            std::array<vk::ShaderStageFlagBits, 2> stages = {
                vk::ShaderStageFlagBits::eVertex,
                vk::ShaderStageFlagBits::eFragment
            };
            std::array<vk::ShaderEXT, 2> shaders = {m_Gpu.forwardVert, m_Gpu.forwardFrag};
            cmd.bindShadersEXT(stages, shaders);

            // Set 0：本帧槽位的 UBO；Set 2：本帧槽位的阴影图
            cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics,
                                   m_Gpu.forwardLayout, 0,
                                   m_Gpu.uboSets[ctx.frame.frameSlot], {});
            cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics,
                                   m_Gpu.forwardLayout, 2,
                                   m_Gpu.shadowSets[ctx.frame.frameSlot], {});

            // ---- 动态状态 ----
            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);
            cmd.setDepthBiasEnable(VK_FALSE);
            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>(ctx.extent.width),
                                  static_cast<float>(ctx.extent.height),
                                  0.0f, 1.0f);
            cmd.setViewportWithCount(viewport);

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

            vk::Bool32 blendEnable = VK_FALSE;
            vk::ColorBlendEquationEXT blendEquation;
            blendEquation.setSrcColorBlendFactor(vk::BlendFactor::eOne)
                         .setDstColorBlendFactor(vk::BlendFactor::eZero)
                         .setColorBlendOp(vk::BlendOp::eAdd)
                         .setSrcAlphaBlendFactor(vk::BlendFactor::eOne)
                         .setDstAlphaBlendFactor(vk::BlendFactor::eZero)
                         .setAlphaBlendOp(vk::BlendOp::eAdd);

            vk::ColorComponentFlags colorWriteMask =
                vk::ColorComponentFlagBits::eR |
                vk::ColorComponentFlagBits::eG |
                vk::ColorComponentFlagBits::eB |
                vk::ColorComponentFlagBits::eA;

            // 两个颜色附件（HDR + velocity）都是无混合直写
            cmd.setColorBlendEnableEXT(0,
                std::array<vk::Bool32, 2>{blendEnable, blendEnable});
            cmd.setColorBlendEquationEXT(0,
                std::array<vk::ColorBlendEquationEXT, 2>{blendEquation, blendEquation});
            cmd.setColorWriteMaskEXT(0,
                std::array<vk::ColorComponentFlags, 2>{colorWriteMask, colorWriteMask});

            // ---- 不透明批次：无混合直写，深度测试 + 写 ----
            for (const auto& batch : m_Frame.mainView.batches) {
                cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics,
                                       m_Gpu.forwardLayout, 1,
                                       batch.mesh->GetMaterial()->GetDescriptorSet(), {});
                batch.mesh->Draw(cmd, batch.count, batch.first, batch.lod);
            }

            // ---- 透明批次：预乘 alpha 混合（线性空间 over），深度只测
            // 不写，远 → 近排序（SceneRenderer 侧已按 batch 平均距离排）。
            // 附件 0（HDR）混合、附件 1（velocity）直写——透明表面的
            // velocity 就是它自身的运动，覆盖自己的像素即正确。
            // alpha 语义：frag 输出未预乘线性色 + 覆盖率 alpha，
            // 预乘在 tonemap 段（rgb * a）完成 —— hard constraint。
            if (!m_Frame.translucentView.batches.empty()) {
                vk::ColorBlendEquationEXT premul;
                premul.setSrcColorBlendFactor(vk::BlendFactor::eSrcAlpha)
                      .setDstColorBlendFactor(vk::BlendFactor::eOneMinusSrcAlpha)
                      .setColorBlendOp(vk::BlendOp::eAdd)
                      .setSrcAlphaBlendFactor(vk::BlendFactor::eOne)
                      .setDstAlphaBlendFactor(vk::BlendFactor::eOneMinusSrcAlpha)
                      .setAlphaBlendOp(vk::BlendOp::eAdd);

                cmd.setColorBlendEnableEXT(0,
                    std::array<vk::Bool32, 2>{VK_TRUE, VK_FALSE});
                cmd.setColorBlendEquationEXT(0,
                    std::array<vk::ColorBlendEquationEXT, 2>{premul, blendEquation});
                cmd.setDepthWriteEnableEXT(VK_FALSE);

                for (const auto& batch : m_Frame.translucentView.batches) {
                    cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics,
                                           m_Gpu.forwardLayout, 1,
                                           batch.mesh->GetMaterial()->GetDescriptorSet(), {});
                    batch.mesh->Draw(cmd, batch.count, batch.first, batch.lod);
                }

                cmd.setDepthWriteEnableEXT(VK_TRUE);   // 恢复（下帧 opaque 段也会重设）
            }

            cmd.endRendering();
        })
        .Write(m_HDRColor, RG::ResourceUse::ColorAttachment)
        .Write(m_Velocity, RG::ResourceUse::ColorAttachment)
        .Write(m_Depth, RG::ResourceUse::DepthAttachment)
        .Read(m_ShadowMaps[0], RG::ResourceUse::ShaderRead,
              vk::ShaderStageFlagBits::eFragment)
        .Read(m_ShadowMaps[1], RG::ResourceUse::ShaderRead,
              vk::ShaderStageFlagBits::eFragment)
        .Read(m_ShadowMaps[2], RG::ResourceUse::ShaderRead,
              vk::ShaderStageFlagBits::eFragment);
    }

} // namespace Engine::Passes