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

#include <Engine/Renderer/Passes/BloomPass.h>
#include <Engine/Core/FrameContext.h>

namespace Engine::Passes
{
    BloomPass::BloomPass(FullscreenGpu& gpu, PostProcessParams& params,
                         RG::TextureHandle hdrColor,
                         RG::TextureHandle bloomA, RG::TextureHandle bloomB)
        : m_Gpu(gpu), m_Params(params),
          m_HDR(hdrColor), m_A(bloomA), m_B(bloomB)
    {
        m_ExtractSets.resize(kMaxFramesInFlight);
        m_BlurHSets.resize(kMaxFramesInFlight);
        m_BlurVSets.resize(kMaxFramesInFlight);
        for (u32 slot = 0; slot < kMaxFramesInFlight; ++slot) {
            m_ExtractSets[slot] = m_Gpu.AllocateSet();
            m_BlurHSets[slot] = m_Gpu.AllocateSet();
            m_BlurVSets[slot] = m_Gpu.AllocateSet();
        }
    }

    void BloomPass::Register(RG::RenderGraph& rg)
    {
        // Extract：HDR → 半分辨率 A（软阈值由 frag 的 param0 传入）
        rg.AddPass("BloomExtract",
            [this](RG::RenderGraphContext& ctx) {
                m_Gpu.RecordFullscreen(ctx, ctx.View(m_A), m_Gpu.extractFrag,
                                       m_ExtractSets[ctx.frame.frameSlot],
                                       FullscreenPush{0.0f, m_Params.bloomThreshold, 0.0f, 0.0f});
            })
            .Read(m_HDR, RG::ResourceUse::ShaderRead,
                  vk::ShaderStageFlagBits::eFragment)
            .Write(m_A, RG::ResourceUse::ColorAttachment);

        // BlurH：A → B（水平步长经 push 传入）
        rg.AddPass("BloomBlurH",
            [this](RG::RenderGraphContext& ctx) {
                m_Gpu.RecordFullscreen(ctx, ctx.View(m_B), m_Gpu.blurFrag,
                                       m_BlurHSets[ctx.frame.frameSlot],
                                       FullscreenPush{0.0f,
                                                      1.0f / ctx.extent.width, 0.0f,
                                                      0.0f});
            })
            .Read(m_A, RG::ResourceUse::ShaderRead,
                  vk::ShaderStageFlagBits::eFragment)
            .Write(m_B, RG::ResourceUse::ColorAttachment);

        // BlurV：B → A（垂直步长）
        rg.AddPass("BloomBlurV",
            [this](RG::RenderGraphContext& ctx) {
                m_Gpu.RecordFullscreen(ctx, ctx.View(m_A), m_Gpu.blurFrag,
                                       m_BlurVSets[ctx.frame.frameSlot],
                                       FullscreenPush{0.0f, 0.0f,
                                                      1.0f / ctx.extent.height,
                                                      0.0f});
            })
            .Read(m_B, RG::ResourceUse::ShaderRead,
                  vk::ShaderStageFlagBits::eFragment)
            .Write(m_A, RG::ResourceUse::ColorAttachment);
    }

    void BloomPass::OnTargetsRebuilt(RG::RenderGraph& rg)
    {
        for (u32 slot = 0; slot < kMaxFramesInFlight; ++slot) {
            m_Gpu.UpdateSetInput(m_ExtractSets[slot], 0, rg, m_HDR, slot);
            m_Gpu.UpdateSetInput(m_BlurHSets[slot], 0, rg, m_A, slot);
            m_Gpu.UpdateSetInput(m_BlurVSets[slot], 0, rg, m_B, slot);
        }
    }

} // namespace Engine::Passes