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

#include <Engine/Renderer/Passes/SceneGpu.h>
#include <Engine/Renderer/Material.h>
#include <Engine/Renderer/Mesh.h>      // IndirectGroup 取数（inline 经 Mesh）
#include <Engine/Renderer/Cubemap.h>
#include <Engine/Vulkan/Device.h>
#include <Engine/Vulkan/Buffer.h>
#include <Engine/Vulkan/DescriptorManager.h>
#include <Engine/Vulkan/ShaderManager.h>
#include <Engine/Vulkan/UniformBuffer.h>
#include <Engine/Core/Log.h>

#include <glm/glm.hpp>

namespace Engine::Passes
{
    SceneGpu::SceneGpu(Device& device, MemoryAllocator& allocator,
                       DescriptorManager& descriptors,
                       ShaderManager& shaders,
                       UniformBuffer& ubo, UniformBuffer& objectBuffer,
                       UniformBuffer& prevObjectBuffer,
                       UniformBuffer& skinOffsetBuffer,
                       UniformBuffer& skinPaletteBuffer)
        : m_Device(device), m_Allocator(allocator), m_Descriptors(descriptors)
    {
        // Set 0: UBO + GPUDriven 实例数据（SSBO，vertex 读取）
        // + 上一帧实例矩阵（TAA velocity，vertex 读取）
        // + 骨骼蒙皮：SkinOffsets（实例→palette 基址）+ SkinPalettes（骨骼矩阵）
        std::vector<DescriptorBinding> uboBindings = {
            {0, vk::DescriptorType::eUniformBuffer,
             vk::ShaderStageFlagBits::eVertex | vk::ShaderStageFlagBits::eFragment},
            {1, vk::DescriptorType::eStorageBuffer,
             vk::ShaderStageFlagBits::eVertex},
            {2, vk::DescriptorType::eStorageBuffer,
             vk::ShaderStageFlagBits::eVertex},
            {3, vk::DescriptorType::eStorageBuffer,
             vk::ShaderStageFlagBits::eVertex},
            {4, vk::DescriptorType::eStorageBuffer,
             vk::ShaderStageFlagBits::eVertex},
        };
        uboSetLayout = descriptors.CreateLayout(uboBindings);

        // Set 1: PBR 材质纹理槽（binding = Material::Slot：0 albedo /
        // 1 normal / 2 metallicRoughness / 3 emissive）+ 参数 UBO（binding 4）
        std::vector<DescriptorBinding> texBindings;
        for (u32 slot = 0; slot < Material::kSlotCount; ++slot) {
            texBindings.push_back(
                {slot, vk::DescriptorType::eCombinedImageSampler,
                 vk::ShaderStageFlagBits::eFragment});
        }
        texBindings.push_back(
            {Material::kParamsBinding, vk::DescriptorType::eUniformBuffer,
             vk::ShaderStageFlagBits::eFragment});
        textureSetLayout = descriptors.CreateLayout(texBindings);

        // Set 2: CSM 级联阴影图（0/1/2）+ IBL 环境三件套（3 irradiance
        // / 4 pre-filtered / 5 BRDF LUT），均 fragment 采样
        std::vector<DescriptorBinding> shadowBindings;
        for (u32 c = 0; c < kMaxCascades; ++c) {
            shadowBindings.push_back(
                {c, vk::DescriptorType::eCombinedImageSampler,
                 vk::ShaderStageFlagBits::eFragment});
        }
        for (const u32 b : {3u, 4u, 5u}) {
            shadowBindings.push_back(
                {b, vk::DescriptorType::eCombinedImageSampler,
                 vk::ShaderStageFlagBits::eFragment});
        }
        shadowSetLayout = descriptors.CreateLayout(shadowBindings);

        // Forward pipeline layout：三个 set；无 push（GPUDriven 经 SSBO）
        forwardLayout = descriptors.CreatePipelineLayout(
            {uboSetLayout, textureSetLayout, shadowSetLayout}, 0);

        // Shadow pipeline layout：Set0 = UBO/ObjectBuffer（实例矩阵），
        // push = lightVP 64B
        shadowLayout = descriptors.CreatePipelineLayout(
            std::vector<vk::DescriptorSetLayout>{uboSetLayout},
            sizeof(glm::mat4));

        const std::string shaderDir = "shaders/";
        // forward vert/frag 均无 push（range 定义一致：都没有，VUID-08878）
        auto vert = shaders.CreateFromFile(
            shaderDir + "simple_shader.vert.spv",
            vk::ShaderStageFlagBits::eVertex,
            {uboSetLayout, textureSetLayout, shadowSetLayout}, 0);
        auto frag = shaders.CreateFromFile(
            shaderDir + "simple_shader.frag.spv",
            vk::ShaderStageFlagBits::eFragment,
            {uboSetLayout, textureSetLayout, shadowSetLayout}, 0);
        // 透明变体：同一套布局，frag 末尾写 gl_FragDepth=1.0
        auto transFrag = shaders.CreateFromFile(
            shaderDir + "simple_shader_translucent.frag.spv",
            vk::ShaderStageFlagBits::eFragment,
            {uboSetLayout, textureSetLayout, shadowSetLayout}, 0);

        // shadow shader：depth-only，push = lightVP
        auto shadowVertShader = shaders.CreateFromFile(
            shaderDir + "shadow.vert.spv",
            vk::ShaderStageFlagBits::eVertex,
            std::vector<vk::DescriptorSetLayout>{uboSetLayout},
            sizeof(glm::mat4));
        auto shadowFragShader = shaders.CreateFromFile(
            shaderDir + "shadow.frag.spv",
            vk::ShaderStageFlagBits::eFragment,
            std::vector<vk::DescriptorSetLayout>{uboSetLayout},
            sizeof(glm::mat4));

        forwardVert = vert.Handle;
        forwardFrag = frag.Handle;
        translucentFrag = transFrag.Handle;
        shadowVert = shadowVertShader.Handle;
        shadowFrag = shadowFragShader.Handle;

        if (!forwardVert || !forwardFrag || !translucentFrag ||
            !shadowVert || !shadowFrag) {
            ENGINE_FATAL("SceneGpu: failed to create shader objects");
        }

        // 阴影采样器：clamp-to-border 白色 = 光视锥外自动无阴影
        //（省 shader 出界分支）；开启硬件深度比较（sampler2DShadow）
        // ——单次采样即 2×2 PCF，配合 shader 3×3 tap 得软阴影
        vk::SamplerCreateInfo shadowSamplerInfo;
        shadowSamplerInfo.setMagFilter(vk::Filter::eLinear)
                         .setMinFilter(vk::Filter::eLinear)
                         .setMipmapMode(vk::SamplerMipmapMode::eNearest)
                         .setAddressModeU(vk::SamplerAddressMode::eClampToBorder)
                         .setAddressModeV(vk::SamplerAddressMode::eClampToBorder)
                         .setAddressModeW(vk::SamplerAddressMode::eClampToBorder)
                         .setBorderColor(vk::BorderColor::eFloatOpaqueWhite)
                         .setUnnormalizedCoordinates(VK_FALSE)
                         .setCompareEnable(VK_TRUE)
                         // ref < stored → lit=1：片元比遮挡物近即被照亮
                         .setCompareOp(vk::CompareOp::eLess);
        shadowSampler = device.GetDevice().createSampler(shadowSamplerInfo);

        // per-slot 描述符分配 + Set0 绑定（buffer 句柄全程稳定，不用重绑）
        uboSets.resize(kMaxFramesInFlight);
        shadowSets.resize(kMaxFramesInFlight);
        for (u32 slot = 0; slot < kMaxFramesInFlight; ++slot) {
            uboSets[slot] = descriptors.AllocateDescriptorSet(uboSetLayout);
            descriptors.UpdateUniformBuffer(
                uboSets[slot], 0, ubo.Get(slot), ubo.GetSize());
            descriptors.UpdateStorageBuffer(
                uboSets[slot], 1, objectBuffer.Get(slot), objectBuffer.GetSize());
            descriptors.UpdateStorageBuffer(
                uboSets[slot], 2, prevObjectBuffer.Get(slot),
                prevObjectBuffer.GetSize());
            descriptors.UpdateStorageBuffer(
                uboSets[slot], 3, skinOffsetBuffer.Get(slot),
                skinOffsetBuffer.GetSize());
            descriptors.UpdateStorageBuffer(
                uboSets[slot], 4, skinPaletteBuffer.Get(slot),
                skinPaletteBuffer.GetSize());

            shadowSets[slot] = descriptors.AllocateDescriptorSet(shadowSetLayout);
        }
        // shadowSets 的图片绑定在 Compile 后由 BindShadowMap 完成

        // ---- MDI per-slot ring：命令 + count（CPU 直写，host visible）----
        // 布局 [main | casc0 | casc1 | casc2] × kMaxInstances 命令预算；
        // 两者 usage 都要 eIndirectBuffer（count buffer 的 spec 要求）
        const vk::DeviceSize cmdBytes = vk::DeviceSize(kIndirectViewMain + 1 + kMaxCascades) *
                                        kMaxInstances * sizeof(IndirectDrawCmd);
        const vk::DeviceSize cntBytes = vk::DeviceSize(kIndirectViewMain + 1 + kMaxCascades) *
                                        kMaxInstances * sizeof(u32);
        constexpr auto memProps = vk::MemoryPropertyFlagBits::eHostVisible |
                                  vk::MemoryPropertyFlagBits::eHostCoherent;
        for (u32 slot = 0; slot < kMaxFramesInFlight; ++slot) {
            indirectCmdBuffers[slot] = CreateScope<Buffer>(
                device, allocator, cmdBytes,
                vk::BufferUsageFlagBits::eIndirectBuffer, memProps);
            indirectCountBuffers[slot] = CreateScope<Buffer>(
                device, allocator, cntBytes,
                vk::BufferUsageFlagBits::eIndirectBuffer, memProps);
        }
    }

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

        if (forwardVert) vkDevice.destroyShaderEXT(forwardVert);
        if (forwardFrag) vkDevice.destroyShaderEXT(forwardFrag);
        if (translucentFrag) vkDevice.destroyShaderEXT(translucentFrag);
        if (shadowVert) vkDevice.destroyShaderEXT(shadowVert);
        if (shadowFrag) vkDevice.destroyShaderEXT(shadowFrag);

        vkDevice.destroySampler(shadowSampler);
        vkDevice.destroyPipelineLayout(forwardLayout);
        vkDevice.destroyPipelineLayout(shadowLayout);
        vkDevice.destroyDescriptorSetLayout(uboSetLayout);
        vkDevice.destroyDescriptorSetLayout(textureSetLayout);
        vkDevice.destroyDescriptorSetLayout(shadowSetLayout);
    }

    void SceneGpu::BindShadowMaps(RG::RenderGraph& rg,
                                  const std::array<RG::TextureHandle, kMaxCascades>& maps)
    {
        // 物理镜像按帧槽位各一份：每个槽位的描述符绑各自的 view
        for (u32 slot = 0; slot < kMaxFramesInFlight; ++slot) {
            for (u32 c = 0; c < kMaxCascades; ++c) {
                m_Descriptors.UpdateTexture(
                    shadowSets[slot], c, shadowSampler,
                    rg.GetPhysicalView(maps[c], slot));
            }
        }
    }

    // ---- MDI：per-slot 上传 + 组绘制 ----

    void SceneGpu::UploadIndirect(u32 frameSlot, const FrameSceneData& frame)
    {
        auto uploadView = [&](u32 viewIdx, const CullResult& v) {
            if (v.indirectCommands.empty()) return;
            const vk::DeviceSize base =
                vk::DeviceSize(viewIdx) * kMaxInstances;
            indirectCmdBuffers[frameSlot]->Upload(
                v.indirectCommands.data(),
                v.indirectCommands.size() * sizeof(IndirectDrawCmd),
                base * sizeof(IndirectDrawCmd));
            // 每组一个 count，下标 = 组首命令下标（drawIndexedIndirectCount
            // 从此读取实际执行条数——GPU 剔除切片接管时只换写入者）
            for (const auto& g : v.indirectGroups) {
                indirectCountBuffers[frameSlot]->Upload(
                    &g.commandCount, sizeof(u32),
                    (base + g.firstCommand) * sizeof(u32));
            }
        };
        uploadView(kIndirectViewMain, frame.mainView);
        for (u32 c = 0; c < kMaxCascades; ++c) {
            uploadView(kIndirectViewMain + 1 + c, frame.shadowCascades[c]);
        }
    }

    void SceneGpu::DrawIndirectGroup(vk::CommandBuffer cmd, u32 frameSlot,
                                     u32 viewIndex, const IndirectGroup& g)
    {
        // 组内全部命令共享同一 megabuffer：bind 一次，逐 mesh 定位由
        // 命令里的 firstIndex/vertexOffset 承担
        cmd.bindVertexBuffers(0, g.mesh->GetVertexBuffer(), {vk::DeviceSize(0)});
        cmd.bindIndexBuffer(g.mesh->GetIndexBuffer(), vk::DeviceSize(0),
                            vk::IndexType::eUint32);
        const vk::DeviceSize base =
            vk::DeviceSize(viewIndex) * kMaxInstances;

        cmd.drawIndexedIndirectCount(
            indirectCmdBuffers[frameSlot]->Get(),
            base * sizeof(IndirectDrawCmd),
            indirectCountBuffers[frameSlot]->Get(),
            (base + g.firstCommand) * sizeof(u32),
            g.commandCount,                 // maxDrawCount = CPU 打包的组大小
            sizeof(IndirectDrawCmd));
    }

    void SceneGpu::BindEnvironment(const Cubemap& env)
    {
        // IBL 资源 view 全程稳定（不随 RG 重编译变化），绑一次即可。
        // cube 资源必须用带 mip 的环境采样器（textureLod 显式 lod 会被
        // clamp 到 maxLod；shadowSampler maxLod=0 会让 pre-filter 永远
        // 停在最锐的 mip）
        for (u32 slot = 0; slot < kMaxFramesInFlight; ++slot) {
            m_Descriptors.UpdateTexture(shadowSets[slot], 3, env.GetSampler(),
                                        env.GetIrradianceView());
            m_Descriptors.UpdateTexture(shadowSets[slot], 4, env.GetSampler(),
                                        env.GetPrefilteredView());
            m_Descriptors.UpdateTexture(shadowSets[slot], 5,
                                        env.GetBrdfLutSampler(),
                                        env.GetBrdfLutView());
        }
    }

} // namespace Engine::Passes