/**
 * @file shader_resource_manager.cpp
 * @brief 着色器资源管理器实现
 */

#include "shader_resource_manager.h"

#include <algorithm>
#include <ranges>

namespace mirage {

// ============================================================================
// 构造与析构
// ============================================================================

shader_resource_manager::shader_resource_manager(
    const logical_device& device,
    resource_manager&     res_mgr
) : device_(device.get_handle()), res_mgr_(res_mgr) {
}

shader_resource_manager::~shader_resource_manager() {
    cleanup();
}

// ============================================================================
// 初始化
// ============================================================================

void shader_resource_manager::initialize(vk::RenderPass render_pass, uint32_t frames_in_flight) {
    render_pass_      = render_pass;
    frames_in_flight_ = frames_in_flight;

    // 获取 Uniform Buffer 对齐要求
    VkPhysicalDeviceProperties properties;
    vkGetPhysicalDeviceProperties(
        res_mgr_.get_physical_device(),
        &properties
    );
    min_uniform_buffer_offset_alignment_ = properties.limits.minUniformBufferOffsetAlignment;

    // 创建共享资源
    create_pipeline_layout();
    create_sampler();

    // 创建每帧的描述符池和 Uniform Buffer
    create_descriptor_pools(frames_in_flight);
}

void shader_resource_manager::create_descriptor_pools(uint32_t frames_in_flight) {
    frame_resources_.resize(frames_in_flight);

    // 获取延迟删除队列
    auto* deletion_queue = &res_mgr_.get_deletion_queue();

    // 为每帧创建描述符池
    for (uint32_t i = 0; i < frames_in_flight; ++i) {
        // 每个帧的描述符池配置
        descriptor_pool_config config;
        config.max_sets = 64;  // 每帧最多 64 个描述符集

        // 预估需要的描述符数量
        constexpr uint32_t descriptors_per_frame = 64;
        config.pool_sizes = {
            {vk::DescriptorType::eUniformBuffer, descriptors_per_frame},
            {vk::DescriptorType::eCombinedImageSampler, descriptors_per_frame * 2}
        };

        auto result = res_mgr_.descriptors().create_pool(config);
        if (result) {
            frame_resources_[i].descriptor_pool = std::move(*result);
        }

        frame_resources_[i].deletion_queue = deletion_queue;

        // 创建每帧的 Uniform Buffer
        create_per_frame_uniform_buffer(i, DEFAULT_UNIFORM_BUFFER_CAPACITY);
    }
}

void shader_resource_manager::create_per_frame_uniform_buffer(uint32_t frame_index, size_t capacity) {
    auto& uniform_buf = frame_resources_[frame_index].uniform_buffer;

    // 对齐容量
    size_t aligned_capacity = (capacity + min_uniform_buffer_offset_alignment_ - 1) /
                              min_uniform_buffer_offset_alignment_ *
                              min_uniform_buffer_offset_alignment_;

    auto result = res_mgr_.create_buffer_handle(
        aligned_capacity,
        vk::BufferUsageFlagBits::eUniformBuffer,
        vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent
    );

    if (result) {
        uniform_buf.buffer  = std::move(*result);
        uniform_buf.capacity = aligned_capacity;
        uniform_buf.offset   = 0;
    }
}

void shader_resource_manager::create_pipeline_layout() {
    // 创建共享的 Pipeline Layout，支持完整的 Push Constants（128 字节）
    // 基础部分 80 字节 + 效果部分 48 字节
    vk::PushConstantRange push_constant_range{
        vk::ShaderStageFlagBits::eVertex | vk::ShaderStageFlagBits::eFragment,
        0,    // 偏移
        128   // 大小（基础80 + 效果48，最大保证大小）
    };

    vk::PipelineLayoutCreateInfo create_info{};
    create_info.setPushConstantRanges(push_constant_range);

    auto [result, layout] = device_.createPipelineLayout(create_info);
    if (result == vk::Result::eSuccess) {
        pipeline_layout_ = device_handle_t<vk::PipelineLayout>(
            layout,
            device_deleter<vk::PipelineLayout>{device_}
        );
    }
}

void shader_resource_manager::create_sampler() {
    // 创建通用的纹理采样器
    vk::SamplerCreateInfo create_info{};
    create_info.magFilter    = vk::Filter::eLinear;
    create_info.minFilter    = vk::Filter::eLinear;
    create_info.addressModeU = vk::SamplerAddressMode::eClampToEdge;
    create_info.addressModeV = vk::SamplerAddressMode::eClampToEdge;
    create_info.addressModeW = vk::SamplerAddressMode::eClampToEdge;

    // 使用比较模式进行深度纹理采样
    create_info.compareEnable = VK_FALSE;
    create_info.compareOp     = vk::CompareOp::eNever;

    // 各向异性过滤
    create_info.maxAnisotropy = 1.0f;
    create_info.anisotropyEnable = VK_FALSE;

    // 边框颜色
    create_info.borderColor = vk::BorderColor::eFloatOpaqueWhite;

    // Mip 映射
    create_info.mipmapMode = vk::SamplerMipmapMode::eLinear;
    create_info.mipLodBias = 0.0f;
    create_info.minLod = 0.0f;
    create_info.maxLod = 0.0f;

    auto [result, sampler] = device_.createSampler(create_info);
    if (result == vk::Result::eSuccess) {
        sampler_ = sampler;
    }
}

// ============================================================================
// Pipeline 管理
// ============================================================================

auto shader_resource_manager::get_or_create_pipeline(
    const pipeline_key&                                         key,
    std::span<const uint32_t>                                   vert_spirv,
    std::span<const uint32_t>                                   frag_spirv,
    std::span<const shader_binding_info>                        bindings,
    std::span<const vk::VertexInputBindingDescription>            vertex_bindings,
    std::span<const vk::VertexInputAttributeDescription>          vertex_attributes,
    VkRenderPass                                                render_pass
) -> vk::Pipeline {
    std::lock_guard lock(pipeline_mutex_);

    // 检查缓存
    auto it = pipeline_cache_.find(key);
    if (it != pipeline_cache_.end()) {
        // 更新 LRU
        lru_list_.erase(it->second.lru_iter_);
        it->second.lru_iter_ = lru_list_.insert(lru_list_.end(), key);
        it->second.last_used_frame = current_frame_index_;
        return it->second.resources.pipeline.get();
    }

    // 创建 Descriptor Set Layout
    auto desc_layout = get_or_create_descriptor_set_layout(bindings);

    // 创建 Shader Modules
    auto vert_module = create_shader_module(vert_spirv);
    auto frag_module = create_shader_module(frag_spirv);

    if (!vert_module || !frag_module) {
        return VK_NULL_HANDLE;
    }

    // 创建 Pipeline
    auto pipeline_res = create_pipeline(
        std::move(vert_module),
        std::move(frag_module),
        desc_layout,
        vertex_bindings,
        vertex_attributes,
        render_pass
    );

    if (!pipeline_res) {
        return VK_NULL_HANDLE;
    }

    // LRU 淘汰
    if (pipeline_cache_.size() >= MAX_PIPELINE_CACHE_SIZE) {
        // 移除最旧的 Pipeline
        auto oldest_key = lru_list_.front();
        lru_list_.pop_front();

        auto old_it = pipeline_cache_.find(oldest_key);
        if (old_it != pipeline_cache_.end()) {
            // 延迟删除
            auto& frame_res = frame_resources_[current_frame_index_ % frames_in_flight_];
            if (frame_res.deletion_queue) {
                old_it->second.resources.schedule_deletion(
                    *frame_res.deletion_queue,
                    current_frame_index_
                );
            } else {
                old_it->second.resources.reset();
            }
            pipeline_cache_.erase(old_it);
        }
    }

    // 插入新 Pipeline
    auto lru_iter = lru_list_.insert(lru_list_.end(), key);
    pipeline_cache_[key] = pipeline_cache_entry{
        .resources    = std::move(pipeline_res),
        .last_used_frame = current_frame_index_,
        .lru_iter_    = lru_iter
    };

    return pipeline_cache_[key].resources.pipeline.get();
}

void shader_resource_manager::cleanup_unused_pipelines(
    uint64_t current_frame,
    uint32_t  max_unused_frames
) {
    std::lock_guard lock(pipeline_mutex_);

    auto it = pipeline_cache_.begin();
    while (it != pipeline_cache_.end()) {
        if (current_frame - it->second.last_used_frame > max_unused_frames) {
            // 延迟删除
            auto& frame_res = frame_resources_[current_frame % frames_in_flight_];
            if (frame_res.deletion_queue) {
                it->second.resources.schedule_deletion(
                    *frame_res.deletion_queue,
                    current_frame
                );
            } else {
                it->second.resources.reset();
            }

            lru_list_.erase(it->second.lru_iter_);
            it = pipeline_cache_.erase(it);
        } else {
            ++it;
        }
    }
}

auto shader_resource_manager::create_shader_module(
    std::span<const uint32_t> spirv_data
) -> device_handle_t<vk::ShaderModule> {
    vk::ShaderModuleCreateInfo create_info{};
    create_info.codeSize = spirv_data.size_bytes();
    create_info.pCode    = spirv_data.data();

    auto [result, module] = device_.createShaderModule(create_info);
    if (result != vk::Result::eSuccess) {
        return {};
    }

    return device_handle_t<vk::ShaderModule>(
        module,
        device_deleter<vk::ShaderModule>{device_}
    );
}

auto shader_resource_manager::create_pipeline(
    device_handle_t<vk::ShaderModule>                           vert_module,
    device_handle_t<vk::ShaderModule>                           frag_module,
    vk::DescriptorSetLayout                                     desc_layout,
    std::span<const vk::VertexInputBindingDescription>            vertex_bindings,
    std::span<const vk::VertexInputAttributeDescription>          vertex_attributes,
    VkRenderPass                                                render_pass
) -> pipeline_resources_t {
    pipeline_resources_t resources;

    // Shader Stage
    vk::PipelineShaderStageCreateInfo vertex_stage_info{};
    vertex_stage_info.setModule(vert_module);
    vertex_stage_info.setStage(vk::ShaderStageFlagBits::eVertex);
    vertex_stage_info.setPName("main");
    vk::PipelineShaderStageCreateInfo fragment_stage_info{};
    fragment_stage_info.setModule(frag_module);
    fragment_stage_info.setStage(vk::ShaderStageFlagBits::eFragment);
    fragment_stage_info.setPName("main");

    std::array shader_stages = {vertex_stage_info, fragment_stage_info};

    // Vertex Input
    vk::PipelineVertexInputStateCreateInfo vertex_input_info{};
    vertex_input_info.setVertexBindingDescriptions(vertex_bindings);
    vertex_input_info.setVertexAttributeDescriptions(vertex_attributes);

    // Input Assembly
    vk::PipelineInputAssemblyStateCreateInfo input_assembly_info{};
    input_assembly_info.topology = vk::PrimitiveTopology::eTriangleStrip;

    // Viewport
    vk::Viewport viewport{};
    viewport.width    = 1.0f;
    viewport.height   = 1.0f;
    viewport.minDepth = 0.0f;
    viewport.maxDepth = 1.0f;

    vk::Rect2D scissor{};
    scissor.extent.setHeight(1);
    scissor.extent.setWidth(1);

    vk::PipelineViewportStateCreateInfo viewport_info{};
    viewport_info.setViewports(viewport);
    viewport_info.setScissors(scissor);

    // Rasterizer
    vk::PipelineRasterizationStateCreateInfo rasterization_info{};
    rasterization_info.polygonMode             = vk::PolygonMode::eFill;
    rasterization_info.cullMode                = vk::CullModeFlagBits::eNone;
    rasterization_info.frontFace               = vk::FrontFace::eClockwise;
    rasterization_info.depthBiasEnable         = VK_FALSE;
    rasterization_info.rasterizerDiscardEnable = VK_FALSE;
    rasterization_info.lineWidth               = 1.0f;

    // Multisampling
    vk::PipelineMultisampleStateCreateInfo multisample_info{};
    multisample_info.rasterizationSamples = vk::SampleCountFlagBits::e1;

    // Depth Stencil
    vk::PipelineDepthStencilStateCreateInfo depth_stencil_info{};
    depth_stencil_info.depthTestEnable  = VK_FALSE;
    depth_stencil_info.depthWriteEnable = VK_FALSE;
    depth_stencil_info.stencilTestEnable = VK_FALSE;

    // Color Blend
    vk::PipelineColorBlendAttachmentState color_blend_attachment{};
    color_blend_attachment.colorWriteMask = vk::ColorComponentFlagBits::eR |
                                            vk::ColorComponentFlagBits::eG |
                                            vk::ColorComponentFlagBits::eB |
                                            vk::ColorComponentFlagBits::eA;
    color_blend_attachment.blendEnable    = VK_FALSE;

    vk::PipelineColorBlendStateCreateInfo color_blend_info{};
    color_blend_info.setAttachments(color_blend_attachment);

    // Pipeline Layout（支持完整的 Push Constants）
    vk::PushConstantRange push_constant_range{
        vk::ShaderStageFlagBits::eVertex | vk::ShaderStageFlagBits::eFragment,
        0,    // 偏移
        128   // 大小（基础80 + 效果48，最大保证大小）
    };

    vk::PipelineLayoutCreateInfo pipeline_layout_info{};
    pipeline_layout_info.setSetLayouts(desc_layout);
    pipeline_layout_info.setPushConstantRanges(push_constant_range);

    auto [layout_result, pipeline_layout] = device_.createPipelineLayout(pipeline_layout_info);
    if (layout_result != vk::Result::eSuccess) {
        return {};
    }

    // Graphics Pipeline
    vk::GraphicsPipelineCreateInfo pipeline_info{};
    pipeline_info.setStages(shader_stages);
    pipeline_info.setPVertexInputState(&vertex_input_info);
    pipeline_info.setPInputAssemblyState(&input_assembly_info);
    pipeline_info.setPViewportState(&viewport_info);
    pipeline_info.setPRasterizationState(&rasterization_info);
    pipeline_info.setPMultisampleState(&multisample_info);
    pipeline_info.setPDepthStencilState(&depth_stencil_info);
    pipeline_info.setPColorBlendState(&color_blend_info);
    pipeline_info.layout     = pipeline_layout;
    pipeline_info.renderPass = render_pass;
    pipeline_info.subpass    = 0;

    auto [pipeline_result, pipeline] = device_.createGraphicsPipeline(
        VK_NULL_HANDLE,
        pipeline_info
    );

    if (pipeline_result != vk::Result::eSuccess) {
        device_.destroyPipelineLayout(pipeline_layout);
        return {};
    }

    // 构建资源
    resources.pipeline = device_handle_t<vk::Pipeline>(
        pipeline,
        device_deleter<vk::Pipeline>{device_}
    );
    resources.layout = device_handle_t<vk::PipelineLayout>(
        pipeline_layout,
        device_deleter<vk::PipelineLayout>{device_}
    );
    resources.desc_layout = device_handle_t<vk::DescriptorSetLayout>(
        desc_layout,
        device_deleter<vk::DescriptorSetLayout>{device_}
    );

    return resources;
}

auto shader_resource_manager::compute_vertex_layout_hash(
    std::span<const vk::VertexInputBindingDescription>   bindings,
    std::span<const vk::VertexInputAttributeDescription> attributes
) -> std::size_t {
    std::size_t h = 0;

    for (const auto& b : bindings) {
        h ^= std::hash<uint32_t>{}(b.binding + 0x9e3779b9 + (h << 6) + (h >> 2));
        h ^= std::hash<uint32_t>{}(b.stride + 0x9e3779b9 + (h << 6) + (h >> 2));
        h ^= std::hash<int>{}(static_cast<int>(b.inputRate) + 0x9e3779b9 + (h << 6) + (h >> 2));
    }

    for (const auto& a : attributes) {
        h ^= std::hash<uint32_t>{}(a.location + 0x9e3779b9 + (h << 6) + (h >> 2));
        h ^= std::hash<uint32_t>{}(a.binding + 0x9e3779b9 + (h << 6) + (h >> 2));
        h ^= std::hash<int>{}(static_cast<int>(a.format) + 0x9e3779b9 + (h << 6) + (h >> 2));
        h ^= std::hash<uint32_t>{}(a.offset + 0x9e3779b9 + (h << 6) + (h >> 2));
    }

    return h;
}

// ============================================================================
// Descriptor Set Layout 管理
// ============================================================================

auto shader_resource_manager::get_or_create_descriptor_set_layout(
    std::span<const shader_binding_info> bindings
) -> vk::DescriptorSetLayout {
    // 构建缓存键
    layout_cache_key key;
    key.bindings = std::vector<shader_binding_info>(bindings.begin(), bindings.end());

    std::lock_guard lock(layout_mutex_);

    // 检查缓存
    auto it = layout_cache_.find(key);
    if (it != layout_cache_.end()) {
        return it->second.get();
    }

    // 转换为 Vulkan 绑定描述
    std::vector<vk::DescriptorSetLayoutBinding> vk_bindings;
    vk_bindings.reserve(bindings.size());

    for (const auto& b : bindings) {
        vk::DescriptorSetLayoutBinding layout_binding{};
        layout_binding.setBinding(b.binding);
        layout_binding.setDescriptorType(b.type);
        layout_binding.setDescriptorCount(b.count);
        layout_binding.setStageFlags(b.stage);
        vk_bindings.push_back(layout_binding);
    }

    // 创建 Layout
    vk::DescriptorSetLayoutCreateInfo create_info{};
    create_info.setBindings(vk_bindings);

    auto [result, layout] = device_.createDescriptorSetLayout(create_info);
    if (result != vk::Result::eSuccess) {
        return VK_NULL_HANDLE;
    }

    auto handle = device_handle_t<vk::DescriptorSetLayout>(
        layout,
        device_deleter<vk::DescriptorSetLayout>{device_}
    );

    layout_cache_.emplace(std::move(key), std::move(handle));

    return layout;
}

// ============================================================================
// 描述符集分配
// ============================================================================

auto shader_resource_manager::allocate_descriptor_set(
    vk::DescriptorSetLayout layout
) -> vk::DescriptorSet {
    auto& frame_res = frame_resources_[current_frame_index_ % frames_in_flight_];

    vk::DescriptorSetAllocateInfo alloc_info{};
    alloc_info.descriptorPool = frame_res.descriptor_pool.get();
    alloc_info.descriptorSetCount = 1;
    alloc_info.pSetLayouts        = &layout;

    auto [result, sets] = device_.allocateDescriptorSets(alloc_info);
    if (result != vk::Result::eSuccess) {
        return VK_NULL_HANDLE;
    }

    return sets[0];
}

void shader_resource_manager::update_uniform_buffer_descriptor(
    vk::DescriptorSet   set,
    uint32_t            binding,
    vk::Buffer          buffer,
    VkDeviceSize        offset,
    VkDeviceSize        size
) const {
    vk::DescriptorBufferInfo buffer_info{
        buffer,
        offset,
        size
    };

    vk::WriteDescriptorSet write{};
    write.setDstSet(set);
    write.setDstBinding(binding);
    write.setDescriptorCount(1);
    write.setDescriptorType(vk::DescriptorType::eUniformBuffer);
    write.setPBufferInfo(&buffer_info);

    device_.updateDescriptorSets(write, nullptr);
}

void shader_resource_manager::update_sampler_descriptor(
    vk::DescriptorSet   set,
    uint32_t            binding,
    vk::ImageView       image_view,
    vk::Sampler         sampler
) const {
    vk::DescriptorImageInfo image_info{
        sampler,
        image_view,
        vk::ImageLayout::eShaderReadOnlyOptimal
    };

    vk::WriteDescriptorSet write{};
    write.setDstSet(set);
    write.setDstBinding(binding);
    write.setDescriptorCount(1);
    write.setDescriptorType(vk::DescriptorType::eCombinedImageSampler);
    write.setPImageInfo(&image_info);

    device_.updateDescriptorSets(write, nullptr);
}

// ============================================================================
// Uniform Buffer 分配
// ============================================================================

auto shader_resource_manager::allocate_uniform_buffer(size_t size) -> uniform_buffer_allocation {
    auto& frame_res = frame_resources_[current_frame_index_ % frames_in_flight_];
    auto& uniform_buf = frame_res.uniform_buffer;

    // 对齐大小
    size_t aligned_size = (size + min_uniform_buffer_offset_alignment_ - 1) /
                          min_uniform_buffer_offset_alignment_ *
                          min_uniform_buffer_offset_alignment_;

    // 检查是否有足够空间
    if (uniform_buf.offset + aligned_size > uniform_buf.capacity) {
        // 环形缓冲区回绕（不处理溢出，简单重置）
        uniform_buf.offset = 0;
    }

    uniform_buffer_allocation result{
        .buffer     = uniform_buf.buffer.get(),
        .offset     = uniform_buf.offset,
        .mapped_ptr = uniform_buf.buffer.mapped_data,
        .size       = size
    };

    uniform_buf.offset += aligned_size;

    return result;
}

// ============================================================================
// 顶点 Buffer 分配
// ============================================================================

auto shader_resource_manager::allocate_vertex_buffer(size_t size) -> vk::Buffer {
    auto result = res_mgr_.create_buffer_handle(
        size,
        vk::BufferUsageFlagBits::eVertexBuffer,
        vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent
    );

    if (result) {
        return result->get();
    }

    return VK_NULL_HANDLE;
}

auto shader_resource_manager::allocate_index_buffer(size_t size) -> vk::Buffer {
    auto result = res_mgr_.create_buffer_handle(
        size,
        vk::BufferUsageFlagBits::eIndexBuffer,
        vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent
    );

    if (result) {
        return result->get();
    }

    return VK_NULL_HANDLE;
}

// ============================================================================
// 帧管理
// ============================================================================

void shader_resource_manager::begin_frame(uint32_t frame_index) {
    current_frame_index_ = frame_index;

    // 重置当前帧的描述符池
    auto& frame_res = frame_resources_[frame_index % frames_in_flight_];
    if (frame_res.descriptor_pool) {
        res_mgr_.descriptors().reset_pool(frame_res.descriptor_pool.get());
    }

    // 重置 Uniform Buffer 偏移
    frame_res.uniform_buffer.offset = 0;
}

// ============================================================================
// 清理
// ============================================================================

void shader_resource_manager::cleanup() {
    // 清理 Pipeline 缓存
    {
        std::lock_guard lock(pipeline_mutex_);
        for (auto& entry : pipeline_cache_) {
            entry.second.resources.reset();
        }
        pipeline_cache_.clear();
        lru_list_.clear();
    }

    // 清理 Layout 缓存
    {
        std::lock_guard lock(layout_mutex_);
        layout_cache_.clear();
    }

    // 清理共享资源
    pipeline_layout_.reset();
    if (sampler_) {
        device_.destroySampler(sampler_);
        sampler_ = VK_NULL_HANDLE;
    }

    // 清理每帧资源
    frame_resources_.clear();
    frames_in_flight_ = 0;
}

} // namespace mirage