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

#include <Engine/Vulkan/DescriptorManager.h>
#include <Engine/Vulkan/Device.h>
#include <Engine/Core/Log.h>

namespace Engine
{
    DescriptorManager::DescriptorManager(Device& device)
        : m_Device(device)
    {
        m_Pools.push_back(CreateChunk());
        ENGINE_INFO("DescriptorManager created (growable, chunk = {} sets)",
                    kChunkSets);
    }

    DescriptorManager::~DescriptorManager()
    {
        for (const auto pool : m_Pools) {
            m_Device.GetDevice().destroyDescriptorPool(pool);
        }
    }

    vk::DescriptorPool DescriptorManager::CreateChunk() const
    {
        std::array<vk::DescriptorPoolSize, 4> poolSizes = {
            vk::DescriptorPoolSize(vk::DescriptorType::eUniformBuffer, kChunkSets),
            vk::DescriptorPoolSize(vk::DescriptorType::eCombinedImageSampler, kChunkSets),
            vk::DescriptorPoolSize(vk::DescriptorType::eStorageBuffer, kChunkSets),
            vk::DescriptorPoolSize(vk::DescriptorType::eStorageImage, kChunkSets),
        };

        vk::DescriptorPoolCreateInfo poolInfo;
        poolInfo.setMaxSets(kChunkSets)
                .setPoolSizes(poolSizes)
                .setFlags(vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet);

        return m_Device.GetDevice().createDescriptorPool(poolInfo);
    }

    vk::DescriptorSetLayout DescriptorManager::CreateLayout(const std::vector<DescriptorBinding>& bindings) const
    {
        std::vector<vk::DescriptorSetLayoutBinding> vkBindings;
        for (auto& b : bindings)
        {
            vk::DescriptorSetLayoutBinding vkb;
            vkb.setBinding(b.Binding)
               .setDescriptorType(b.Type)
               .setDescriptorCount(1)
               .setStageFlags(b.Stages);
            vkBindings.push_back(vkb);
        }

        vk::DescriptorSetLayoutCreateInfo layoutInfo;
        layoutInfo.setBindings(vkBindings);

        return m_Device.GetDevice().createDescriptorSetLayout(layoutInfo);
    }

    vk::DescriptorSet DescriptorManager::AllocateDescriptorSet(vk::DescriptorSetLayout layout)
    {
        // 从最新块分配；耗尽（OutOfPool/Fragmented）则追加一块重试
        for (;;) {
            try {
                vk::DescriptorSetAllocateInfo allocInfo;
                allocInfo.setDescriptorPool(m_Pools.back())
                         .setSetLayouts(layout);
                auto sets = m_Device.GetDevice().allocateDescriptorSets(allocInfo);
                return sets[0];
            } catch (const vk::OutOfPoolMemoryError&) {
                m_Pools.push_back(CreateChunk());
                ENGINE_INFO("DescriptorManager: pool chunk grown -> {} chunks",
                            m_Pools.size());
            } catch (const vk::FragmentedPoolError&) {
                m_Pools.push_back(CreateChunk());
                ENGINE_INFO("DescriptorManager: pool fragmented, new chunk -> {}",
                            m_Pools.size());
            }
        }
    }

    void DescriptorManager::UpdateUniformBuffer(const vk::DescriptorSet set, const u32 binding,
                                                const vk::Buffer buffer, const u64 size) const
    {
        vk::DescriptorBufferInfo bufferInfo;
        bufferInfo.setBuffer(buffer).setOffset(0).setRange(size);

        vk::WriteDescriptorSet write;
        write.setDstSet(set)
             .setDstBinding(binding)
             .setDstArrayElement(0)
             .setDescriptorType(vk::DescriptorType::eUniformBuffer)
             .setDescriptorCount(1)
             .setBufferInfo(bufferInfo);

        m_Device.GetDevice().updateDescriptorSets(write, {});
    }

    void DescriptorManager::UpdateStorageBuffer(const vk::DescriptorSet set, const u32 binding,
                                                const vk::Buffer buffer, const u64 size) const
    {
        vk::DescriptorBufferInfo bufferInfo;
        bufferInfo.setBuffer(buffer).setOffset(0).setRange(size);

        vk::WriteDescriptorSet write;
        write.setDstSet(set)
             .setDstBinding(binding)
             .setDstArrayElement(0)
             .setDescriptorType(vk::DescriptorType::eStorageBuffer)
             .setDescriptorCount(1)
             .setBufferInfo(bufferInfo);

        m_Device.GetDevice().updateDescriptorSets(write, {});
    }

    void DescriptorManager::UpdateTexture(const vk::DescriptorSet set, const u32 binding,
                                          const vk::Sampler sampler, const vk::ImageView view) const
    {
        vk::DescriptorImageInfo imageInfo;
        imageInfo.setSampler(sampler)
                 .setImageView(view)
                 .setImageLayout(vk::ImageLayout::eShaderReadOnlyOptimal);

        vk::WriteDescriptorSet write;
        write.setDstSet(set)
             .setDstBinding(binding)
             .setDstArrayElement(0)
             .setDescriptorType(vk::DescriptorType::eCombinedImageSampler)
             .setDescriptorCount(1)
             .setImageInfo(imageInfo);

        m_Device.GetDevice().updateDescriptorSets(write, {});
    }

    void DescriptorManager::UpdateStorageImage(const vk::DescriptorSet set,
                                               const u32 binding,
                                               const vk::ImageView view) const
    {
        // storage image 写（compute 输出）：layout 必须 eGeneral
        vk::DescriptorImageInfo imageInfo;
        imageInfo.setImageView(view)
                 .setImageLayout(vk::ImageLayout::eGeneral);

        vk::WriteDescriptorSet write;
        write.setDstSet(set)
             .setDstBinding(binding)
             .setDstArrayElement(0)
             .setDescriptorType(vk::DescriptorType::eStorageImage)
             .setDescriptorCount(1)
             .setImageInfo(imageInfo);

        m_Device.GetDevice().updateDescriptorSets(write, {});
    }

    vk::PipelineLayout DescriptorManager::CreatePipelineLayout(vk::DescriptorSetLayout layout,
                                                               const u32 pushConstantSize) const
    {
        vk::PipelineLayoutCreateInfo info;
        info.setSetLayouts(layout);

        // 关键：数组声明在 if 外面，生命周期覆盖到 createPipelineLayout
        std::array<vk::PushConstantRange, 1> pushRanges;

        if (pushConstantSize > 0)
        {
            pushRanges[0].setStageFlags(vk::ShaderStageFlagBits::eVertex)
                         .setOffset(0)
                         .setSize(pushConstantSize);
            info.setPushConstantRanges(pushRanges);
        }

        return m_Device.GetDevice().createPipelineLayout(info);
    }

    vk::PipelineLayout DescriptorManager::CreatePipelineLayout(const std::vector<vk::DescriptorSetLayout>& layouts,
                                                               const u32 pushConstantSize,
                                                               const vk::ShaderStageFlags pushStages) const
    {
        vk::PipelineLayoutCreateInfo info;
        info.setSetLayouts(layouts);

        std::array<vk::PushConstantRange, 1> pushRanges;
        if (pushConstantSize > 0) {
            // 与 ShaderManager 一致：程序级声明（默认 vertex|fragment）。
            // 注意 vkCmdPushConstants 的 stageFlags 必须覆盖 range 的
            // 全部 stage（VUID-01796），因此 push 调用点也要带同样的 stage
            pushRanges[0].setStageFlags(pushStages)
                         .setOffset(0)
                         .setSize(pushConstantSize);
            info.setPushConstantRanges(pushRanges);
        }

        return m_Device.GetDevice().createPipelineLayout(info);
    }

} // namespace Engine