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

#include <Engine/Renderer/Passes/FullscreenGpu.h>
#include <Engine/Vulkan/Device.h>
#include <Engine/Vulkan/DescriptorManager.h>
#include <Engine/Vulkan/ShaderManager.h>
#include <Engine/Core/Log.h>

#include <array>

namespace Engine::Passes
{
    FullscreenGpu::FullscreenGpu(Device& device, DescriptorManager& descriptors,
                                 ShaderManager& shaders)
        : m_Device(device), m_Descriptors(descriptors)
    {
        // Set0 = {HDR 颜色图, Bloom 图}，push = 16B 通用结构
        std::vector<DescriptorBinding> bindings = {
            {0, vk::DescriptorType::eCombinedImageSampler,
             vk::ShaderStageFlagBits::eFragment},
            {1, vk::DescriptorType::eCombinedImageSampler,
             vk::ShaderStageFlagBits::eFragment},
        };
        setLayout = descriptors.CreateLayout(bindings);

        layout = descriptors.CreatePipelineLayout(
            std::vector<vk::DescriptorSetLayout>{setLayout},
            static_cast<u32>(sizeof(FullscreenPush)));

        sampler = device.GetDevice().createSampler(
            vk::SamplerCreateInfo{}
                .setMagFilter(vk::Filter::eLinear)
                .setMinFilter(vk::Filter::eLinear)
                .setMipmapMode(vk::SamplerMipmapMode::eNearest)
                .setAddressModeU(vk::SamplerAddressMode::eClampToEdge)
                .setAddressModeV(vk::SamplerAddressMode::eClampToEdge)
                .setAddressModeW(vk::SamplerAddressMode::eClampToEdge));

        const std::string shaderDir = "shaders/";
        auto vertShader = shaders.CreateFromFile(
            shaderDir + "tonemap.vert.spv",
            vk::ShaderStageFlagBits::eVertex,
            std::vector<vk::DescriptorSetLayout>{setLayout},
            static_cast<u32>(sizeof(FullscreenPush)));
        auto tonemapF = shaders.CreateFromFile(
            shaderDir + "tonemap.frag.spv",
            vk::ShaderStageFlagBits::eFragment,
            std::vector<vk::DescriptorSetLayout>{setLayout},
            static_cast<u32>(sizeof(FullscreenPush)));
        auto extractF = shaders.CreateFromFile(
            shaderDir + "bloom_extract.frag.spv",
            vk::ShaderStageFlagBits::eFragment,
            std::vector<vk::DescriptorSetLayout>{setLayout},
            static_cast<u32>(sizeof(FullscreenPush)));
        auto blurF = shaders.CreateFromFile(
            shaderDir + "bloom_blur.frag.spv",
            vk::ShaderStageFlagBits::eFragment,
            std::vector<vk::DescriptorSetLayout>{setLayout},
            static_cast<u32>(sizeof(FullscreenPush)));

        vert = vertShader.Handle;
        tonemapFrag = tonemapF.Handle;
        extractFrag = extractF.Handle;
        blurFrag = blurF.Handle;

        if (!vert || !tonemapFrag || !extractFrag || !blurFrag) {
            ENGINE_FATAL("FullscreenGpu: failed to create shader objects");
        }
    }

    FullscreenGpu::~FullscreenGpu()
    {
        const auto& vkDevice = m_Device.GetDevice();

        if (vert) vkDevice.destroyShaderEXT(vert);
        if (tonemapFrag) vkDevice.destroyShaderEXT(tonemapFrag);
        if (extractFrag) vkDevice.destroyShaderEXT(extractFrag);
        if (blurFrag) vkDevice.destroyShaderEXT(blurFrag);

        vkDevice.destroySampler(sampler);
        vkDevice.destroyPipelineLayout(layout);
        vkDevice.destroyDescriptorSetLayout(setLayout);
    }

    vk::DescriptorSet FullscreenGpu::AllocateSet()
    {
        return m_Descriptors.AllocateDescriptorSet(setLayout);
    }

    void FullscreenGpu::UpdateSetInput(vk::DescriptorSet set, u32 binding,
                                       RG::RenderGraph& rg, RG::TextureHandle input,
                                       u32 slot)
    {
        // 物理镜像按帧槽位各一份：每个槽位的描述符绑各自的 view
        m_Descriptors.UpdateTexture(set, binding, sampler,
                                    rg.GetPhysicalView(input, slot));
    }

    void FullscreenGpu::RecordFullscreen(RG::RenderGraphContext& ctx,
                                         vk::ImageView view,
                                         vk::ShaderEXT frag,
                                         const vk::DescriptorSet& set,
                                         const FullscreenPush& push)
    {
        const vk::CommandBuffer cmd = ctx.cmd;

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

        vk::RenderingAttachmentInfo colorAttachment;
        colorAttachment.setImageView(view)
                       .setImageLayout(vk::ImageLayout::eColorAttachmentOptimal)
                       .setLoadOp(vk::AttachmentLoadOp::eClear)
                       .setStoreOp(vk::AttachmentStoreOp::eStore)
                       .setClearValue(clearValue);

        vk::RenderingInfo renderingInfo;
        renderingInfo.setRenderArea({{0, 0}, ctx.extent})
                     .setLayerCount(1)
                     .setColorAttachments(colorAttachment);

        cmd.beginRendering(renderingInfo);

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

        cmd.bindDescriptorSets(vk::PipelineBindPoint::eGraphics,
                               layout, 0, set, {});

        // ---- 动态状态（无深度附件、无顶点输入）----
        cmd.setRasterizerDiscardEnable(VK_FALSE);
        cmd.setDepthTestEnable(VK_FALSE);
        cmd.setDepthWriteEnable(VK_FALSE);
        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);
        cmd.setVertexInputEXT(
            std::array<vk::VertexInputBindingDescription2EXT, 0>{},
            std::array<vk::VertexInputAttributeDescription2EXT, 0>{});
        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);
        cmd.setScissorWithCount(vk::Rect2D({0, 0}, ctx.extent));

        cmd.pushConstants(layout,
                          vk::ShaderStageFlagBits::eVertex |
                          vk::ShaderStageFlagBits::eFragment,
                          0, sizeof(FullscreenPush), &push);

        cmd.draw(3, 1, 0, 0);
        cmd.endRendering();
    }

} // namespace Engine::Passes