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

#include <Engine/Renderer/Passes/SceneGpu.h>
#include <Engine/Renderer/Material.h>
#include <Engine/Vulkan/Device.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, DescriptorManager& descriptors,
                       ShaderManager& shaders,
                       UniformBuffer& ubo, UniformBuffer& objectBuffer,
                       UniformBuffer& prevObjectBuffer)
        : m_Device(device), m_Descriptors(descriptors)
    {
        // Set 0: UBO + GPUDriven 实例数据（SSBO，vertex 读取）
        // + 上一帧实例矩阵（TAA velocity，vertex 读取）
        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},
        };
        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 级联阴影图（fragment 采样，binding = 级联序号）
        std::vector<DescriptorBinding> shadowBindings;
        for (u32 c = 0; c < kMaxCascades; ++c) {
            shadowBindings.push_back(
                {c, 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);

        // 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;
        shadowVert = shadowVertShader.Handle;
        shadowFrag = shadowFragShader.Handle;

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

        // 阴影采样器：clamp-to-border 白色 = 光视锥之外无阴影
        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_FALSE);
        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());

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

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

        if (forwardVert) vkDevice.destroyShaderEXT(forwardVert);
        if (forwardFrag) vkDevice.destroyShaderEXT(forwardFrag);
        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));
            }
        }
    }

} // namespace Engine::Passes