Extra2D/include/renderer/instance_buffer.h

#pragma once

#include <renderer/rhi/rhi.h>
#include <renderer/rhi/rhi_types.h>
#include <types/math/vec2.h>
#include <types/math/color.h>
#include <types/math/mat4.h>
#include <vector>
#include <memory>

namespace extra2d {

/**
 * @brief 实例数据
 *
 * 单个实例的属性数据，用于实例化渲染
 * 布局遵循 std140 对齐规则
 * 
 * 优化说明：
 * - 旋转使用预计算的 cos/sin 值，避免在着色器中进行三角函数计算
 * - 属性布局优化，便于GPU访问
 */
struct InstanceData {
    // 第一组 16 字节 - 变换数据
    Vec2 position;      // 位置偏移 (8 bytes)
    float rotationCos;  // 旋转角度余弦值 (4 bytes) - 预计算避免着色器三角函数
    float rotationSin;  // 旋转角度正弦值 (4 bytes) - 预计算避免着色器三角函数
    
    // 第二组 16 字节 - 缩放和预留
    Vec2 scale;         // 缩放 (8 bytes)
    float padding1[2];  // 填充到 16 字节对齐 (8 bytes)
    
    // 第三组 16 字节 - 颜色
    Color color;        // 颜色 (16 bytes) - r, g, b, a
    
    // 第四组 16 字节 - UV坐标
    float uvX;          // UV 起始 X (4 bytes)
    float uvY;          // UV 起始 Y (4 bytes)
    float uvWidth;      // UV 宽度 (4 bytes)
    float uvHeight;     // UV 高度 (4 bytes)

    InstanceData()
        : position(0.0f, 0.0f)
        , rotationCos(1.0f)  // cos(0) = 1
        , rotationSin(0.0f)  // sin(0) = 0
        , scale(1.0f, 1.0f)
        , padding1{0.0f, 0.0f}
        , color(Color::White)
        , uvX(0.0f)
        , uvY(0.0f)
        , uvWidth(1.0f)
        , uvHeight(1.0f) {}

    /**
     * @brief 设置旋转角度（自动计算cos/sin）
     * @param angle 旋转角度（弧度）
     */
    void setRotation(float angle) {
        rotationCos = std::cos(angle);
        rotationSin = std::sin(angle);
    }

    /**
     * @brief 获取旋转角度（通过atan2计算，仅用于调试）
     * @return 旋转角度（弧度）
     */
    float getRotation() const {
        return std::atan2(rotationSin, rotationCos);
    }
};

static_assert(sizeof(InstanceData) == 64, "InstanceData size should be 64 bytes for std140 alignment");

/**
 * @brief 脏区域标记
 *
 * 标记需要更新的实例数据区域，支持部分更新GPU缓冲区
 */
struct DirtyRange {
    uint32_t start;     // 起始实例索引
    uint32_t count;     // 实例数量
    
    DirtyRange(uint32_t s = 0, uint32_t c = 0) : start(s), count(c) {}
    
    bool isValid() const { return count > 0; }
    uint32_t end() const { return start + count; }
    
    // 检查是否与另一个区域重叠或相邻
    bool canMerge(const DirtyRange& other) const {
        return (start <= other.end() && other.start <= end()) ||
               (end() + 1 >= other.start && other.end() + 1 >= start);
    }
    
    // 合并两个区域
    void merge(const DirtyRange& other) {
        uint32_t newStart = std::min(start, other.start);
        uint32_t newEnd = std::max(end(), other.end());
        start = newStart;
        count = newEnd - newStart;
    }
};

/**
 * @brief 实例缓冲区
 *
 * 管理实例化渲染的实例数据缓冲区
 * 支持动态更新和双缓冲
 * 
 * 优化特性：
 * - 脏标记系统：只更新变化的实例数据，减少PCIe带宽占用
 * - 批量更新：合并相邻的脏区域，减少GPU上传次数
 */
class InstanceBuffer {
public:
    /**
     * @brief 默认构造函数
     */
    InstanceBuffer();

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

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

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

    /**
     * @brief 初始化实例缓冲区
     * @param maxInstances 最大实例数量
     * @return 初始化是否成功
     */
    bool initialize(uint32_t maxInstances);

    /**
     * @brief 关闭缓冲区
     */
    void shutdown();

    /**
     * @brief 更新实例数据（全量更新，标记整个缓冲区为脏）
     * @param instances 实例数据数组
     * @param count 实例数量
     * @return 更新是否成功
     */
    bool updateInstances(const InstanceData* instances, uint32_t count);

    /**
     * @brief 更新部分实例数据（增量更新）
     * @param instances 实例数据数组
     * @param start 起始实例索引
     * @param count 实例数量
     * @return 更新是否成功
     */
    bool updateInstancesRange(const InstanceData* instances, uint32_t start, uint32_t count);

    /**
     * @brief 添加单个实例
     * @param instance 实例数据
     * @return 实例索引
     */
    uint32_t addInstance(const InstanceData& instance);

    /**
     * @brief 标记实例范围为脏（下次updateGPU时更新）
     * @param start 起始实例索引
     * @param count 实例数量
     */
    void markDirty(uint32_t start, uint32_t count);

    /**
     * @brief 标记整个缓冲区为脏
     */
    void markAllDirty();

    /**
     * @brief 将脏数据上传到GPU
     * @return 上传是否成功
     */
    bool updateGPU();

    /**
     * @brief 清除所有实例
     */
    void clear();

    /**
     * @brief 获取当前实例数量
     * @return 实例数量
     */
    uint32_t getInstanceCount() const { return instanceCount_; }

    /**
     * @brief 获取最大实例数量
     * @return 最大实例数量
     */
    uint32_t getMaxInstances() const { return maxInstances_; }

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

    /**
     * @brief 获取 RHI 缓冲区指针
     * @return 缓冲区指针
     */
    RHIBuffer* getRHIBuffer() const { return bufferHandle_.get(); }

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

    /**
     * @brief 获取脏区域数量
     * @return 脏区域数量
     */
    uint32_t getDirtyRangeCount() const { return static_cast<uint32_t>(dirtyRanges_.size()); }

    /**
     * @brief 检查是否有脏数据
     * @return 是否有脏数据
     */
    bool hasDirtyData() const { return !dirtyRanges_.empty(); }

private:
    BufferHandle bufferHandle_;             // RHI 缓冲区句柄
    uint32_t maxInstances_ = 0;             // 最大实例数量
    uint32_t instanceCount_ = 0;            // 当前实例数量
    std::vector<InstanceData> cpuBuffer_;   // CPU 端缓冲区
    std::vector<DirtyRange> dirtyRanges_;   // 脏区域列表
    
    // 最大脏区域数量，超过则合并为全量更新
    static constexpr uint32_t MAX_DIRTY_RANGES = 8;
    
    /**
     * @brief 添加脏区域，自动合并重叠区域
     * @param range 脏区域
     */
    void addDirtyRange(const DirtyRange& range);
    
    /**
     * @brief 合并所有脏区域为一个
     */
    void mergeAllDirtyRanges();
    
    /**
     * @brief 清空脏区域列表
     */
    void clearDirtyRanges();
};

/**
 * @brief 实例缓冲区管理器
 *
 * 管理多个实例缓冲区的分配和回收
 */
class InstanceBufferManager {
public:
    /**
     * @brief 默认构造函数
     */
    InstanceBufferManager();

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

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

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

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

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

    /**
     * @brief 获取或创建实例缓冲区
     * @param minSize 最小容量（实例数量）
     * @return 实例缓冲区指针
     */
    InstanceBuffer* acquireBuffer(uint32_t minSize);

    /**
     * @brief 回收实例缓冲区
     * @param buffer 缓冲区指针
     */
    void releaseBuffer(InstanceBuffer* buffer);

    /**
     * @brief 重置所有缓冲区
     */
    void reset();

private:
    std::vector<std::unique_ptr<InstanceBuffer>> bufferPool_;
    uint32_t currentBufferIndex_ = 0;
    static constexpr uint32_t DEFAULT_BUFFER_SIZE = 1024;
};

} // namespace extra2d