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

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

namespace Engine::Passes
{
    TAAPass::TAAPass(FullscreenGpu& gpu, FrameSceneData& frame, f32 historyAlpha,
                     RG::TextureHandle hdrColor, RG::TextureHandle velocity,
                     RG::TextureHandle taaColor)
        : m_Gpu(gpu), m_Frame(frame), m_HistoryAlpha(historyAlpha),
          m_HDR(hdrColor), m_Velocity(velocity), m_TAAColor(taaColor)
    {
        m_Sets.resize(kMaxFramesInFlight);
        for (u32 slot = 0; slot < kMaxFramesInFlight; ++slot) {
            m_Sets[slot] = m_Gpu.AllocateSet();
        }
    }

    void TAAPass::Register(RG::RenderGraph& rg)
    {
        rg.AddPass("TAA",
            [this](RG::RenderGraphContext& ctx) {
                // TAA 关闭：alpha=0 → mix(cur, hist, 0) 直通当前帧
                //（抖动已由 SceneRenderer 同步停掉，画面静止无锯齿晃动）
                const f32 alpha = m_Frame.taaEnabled ? m_HistoryAlpha : 0.0f;
                m_Gpu.RecordFullscreen(ctx, ctx.View(m_TAAColor), m_Gpu.taaFrag,
                                       m_Sets[ctx.frame.frameSlot],
                                       FullscreenPush{0.0f, alpha, 0.0f, 0.0f});
            })
            .Read(m_HDR, RG::ResourceUse::ShaderRead,
                  vk::ShaderStageFlagBits::eFragment)
            .Read(m_Velocity, RG::ResourceUse::ShaderRead,
                  vk::ShaderStageFlagBits::eFragment)
            // 自引用历史：TAAColor 的上一帧槽位就是累积历史
            .ReadHistory(m_TAAColor, RG::ResourceUse::ShaderRead,
                         vk::ShaderStageFlagBits::eFragment)
            .Write(m_TAAColor, RG::ResourceUse::ColorAttachment);
    }

    void TAAPass::OnTargetsRebuilt(RG::RenderGraph& rg)
    {
        for (u32 slot = 0; slot < kMaxFramesInFlight; ++slot) {
            m_Gpu.UpdateSetInput(m_Sets[slot], 0, rg, m_HDR, slot);
            m_Gpu.UpdateSetInput(m_Sets[slot], 1, rg, m_Velocity, slot);
            // binding2 = 历史（上一帧槽位镜像）
            m_Gpu.UpdateSetInput(m_Sets[slot], 2, rg, m_TAAColor,
                                 rg.HistorySlot(slot));
        }
    }

} // namespace Engine::Passes