Extra2D/include/renderer/uniform_buffer.h

#pragma once

#include <renderer/rhi/rhi.h>
#include <memory>
#include <vector>

namespace extra2d {

/**
 * @brief 统一缓冲区对象 (UBO)
 *
 * 基于 RHI 的 UBO 包装类，用于高效地传递 uniform 数据到 GPU
 * 支持 std140 布局标准
 */
class UniformBuffer {
public:
    /**
     * @brief 默认构造函数
     */
    UniformBuffer();

    /**
     * @brief 析构函数
     */
    ~UniformBuffer();

    /**
     * @brief 创建 UBO
     * @param size 缓冲区大小（字节）
     * @param binding 绑定槽位
     * @return 创建是否成功
     */
    bool create(uint32_t size, uint32_t binding);

    /**
     * @brief 销毁 UBO
     */
    void destroy();

    /**
     * @brief 更新缓冲区数据
     * @param data 数据源指针
     * @param size 数据大小
     * @param offset 缓冲区偏移（字节）
     */
    void update(const void* data, uint32_t size, uint32_t offset = 0);

    /**
     * @brief 绑定到指定槽位
     * @param binding 绑定槽位
     */
    void bind(uint32_t binding);

    /**
     * @brief 获取缓冲区大小
     * @return 缓冲区大小
     */
    uint32_t getSize() const { return size_; }

    /**
     * @brief 获取 RHI 缓冲区句柄
     * @return RHI 缓冲区句柄
     */
    BufferHandle getHandle() const { return handle_; }

    /**
     * @brief 获取底层 RHI 缓冲区指针
     * @return RHI 缓冲区指针
     */
    RHIBuffer* getRHIBuffer() const;

    /**
     * @brief 检查是否有效
     * @return 是否有效
     */
    bool isValid() const { return handle_.isValid(); }

private:
    BufferHandle handle_;       // RHI 缓冲区句柄
    uint32_t size_ = 0;         // 缓冲区大小
    uint32_t binding_ = 0;      // 绑定槽位
};

/**
 * @brief UBO 管理器
 *
 * 管理多个 UBO 的分配和回收
 * 支持每帧重置和对象池复用
 * 支持双缓冲机制减少 CPU-GPU 等待
 */
class UniformBufferManager {
public:
    /**
     * @brief 默认构造函数
     */
    UniformBufferManager();

    /**
     * @brief 析构函数
     */
    ~UniformBufferManager();

    // 禁止拷贝
    UniformBufferManager(const UniformBufferManager&) = delete;
    UniformBufferManager& operator=(const UniformBufferManager&) = delete;

    // 允许移动
    UniformBufferManager(UniformBufferManager&&) noexcept;
    UniformBufferManager& operator=(UniformBufferManager&&) noexcept;

    /**
     * @brief 初始化管理器
     * @return 初始化是否成功
     */
    bool initialize();

    /**
     * @brief 关闭管理器
     */
    void shutdown();

    /**
     * @brief 获取当前帧的全局 UBO（双缓冲）
     * @param frameIndex 当前帧索引
     * @return 全局 UBO 指针
     */
    UniformBuffer* getGlobalUBO(uint32_t frameIndex);

    /**
     * @brief 获取或创建材质 UBO
     * @param size 缓冲区大小
     * @return UBO 指针
     */
    UniformBuffer* acquireMaterialUBO(uint32_t size);

    /**
     * @brief 重置材质 UBO 池
     *
     * 每帧调用，重置当前索引以便复用 UBO
     */
    void resetMaterialUBOs();

    /**
     * @brief 更新全局 UBO 数据（双缓冲）
     * @param data 数据指针
     * @param size 数据大小
     * @param frameIndex 当前帧索引
     */
    void updateGlobalUBO(const void* data, uint32_t size, uint32_t frameIndex);

    /**
     * @brief 批量更新材质 UBO 数据
     * @param data 材质数据指针
     * @param size 数据大小
     * @param offset 偏移量
     */
    void batchUpdateMaterialUBO(const void* data, uint32_t size, uint32_t offset);

    /**
     * @brief 获取材质 UBO 的 CPU 缓冲区指针（用于批量更新）
     * @return CPU 缓冲区指针
     */
    uint8_t* getMaterialUBOBuffer() { return materialUBOBuffer_.data(); }

    /**
     * @brief 获取材质 UBO 缓冲区大小
     * @return 缓冲区大小
     */
    uint32_t getMaterialUBOBufferSize() const { return static_cast<uint32_t>(materialUBOBuffer_.size()); }

    /**
     * @brief 刷新材质 UBO 数据到 GPU
     */
    void flushMaterialUBO();

private:
    // 全局 UBO 双缓冲（用于 viewProjection、time 等全局数据）
    std::array<std::unique_ptr<UniformBuffer>, 2> globalUBOs_;

    // 材质 UBO 池
    std::vector<std::unique_ptr<UniformBuffer>> materialUBOPool_;
    uint32_t currentUBOIndex_ = 0;

    // 材质 UBO CPU 缓冲区（用于批量更新）
    std::vector<uint8_t> materialUBOBuffer_;
    uint32_t materialUBOBufferOffset_ = 0;
    UniformBuffer* currentMaterialUBO_ = nullptr;

    // 常量
    static constexpr uint32_t INITIAL_UBO_POOL_SIZE = 16;
    static constexpr uint32_t GLOBAL_UBO_SIZE = 256;      // 足够存储 viewProjection + time + screenSize
    static constexpr uint32_t GLOBAL_UBO_BINDING = 0;     // 全局 UBO 绑定槽位
    static constexpr uint32_t MATERIAL_UBO_BINDING = 1;   // 材质 UBO 绑定槽位
    static constexpr uint32_t MATERIAL_UBO_BUFFER_SIZE = 1024 * 1024; // 1MB 材质 UBO 缓冲区
};

} // namespace extra2d