Extra2D/include/renderer/command_queue.h

#pragma once

#include <array>
#include <memory>
#include <renderer/rhi/rhi.h>
#include <renderer/rhi/rhi_command_list.h>
#include <renderer/uniform_buffer.h>
#include <types/math/color.h>
#include <types/math/transform.h>
#include <unordered_map>
#include <vector>

namespace extra2d {

// 前向声明
class CommandQueue;
class Material;
class Mesh;
class UniformBufferManager;
template <typename T> class Ptr;

/**
 * @brief 绘制命令键
 *
 * 64位排序键，用于高效排序绘制命令
 * 布局：| 材质ID (32位) | 深度 (16位) | 层 (8位) | 标志 (8位) |
 */
struct DrawKey {
  uint64_t value;

  DrawKey() : value(0) {}
  explicit DrawKey(uint64_t v) : value(v) {}

  // 构建排序键
  static DrawKey make(uint32_t materialId, uint16_t depth, uint8_t layer,
                      uint8_t flags = 0) {
    DrawKey key;
    key.value = (static_cast<uint64_t>(materialId) << 32) |
                (static_cast<uint64_t>(depth) << 16) |
                (static_cast<uint64_t>(layer) << 8) |
                static_cast<uint64_t>(flags);
    return key;
  }

  uint32_t getMaterialId() const { return static_cast<uint32_t>(value >> 32); }
  uint16_t getDepth() const {
    return static_cast<uint16_t>((value >> 16) & 0xFFFF);
  }
  uint8_t getLayer() const { return static_cast<uint8_t>((value >> 8) & 0xFF); }
  uint8_t getFlags() const { return static_cast<uint8_t>(value & 0xFF); }

  bool operator<(const DrawKey &other) const { return value < other.value; }
  bool operator>(const DrawKey &other) const { return value > other.value; }
};

/**
 * @brief 绘制命令
 *
 * 封装一次绘制调用的所有信息
 */
struct DrawCommand {
  DrawKey key;                           // 排序键
  PipelineHandle pipeline;               // 管线状态
  BufferHandle vertexBuffer;             // 顶点缓冲区
  BufferHandle indexBuffer;              // 索引缓冲区（可选）
  uint32_t vertexCount;                  // 顶点数量
  uint32_t indexCount;                   // 索引数量
  uint32_t instanceCount;                // 实例数量（1表示非实例化）
  std::array<TextureHandle, 8> textures; // 纹理数组
  uint32_t textureCount;                 // 纹理数量
  BufferHandle materialUBO;              // 材质 UBO
  uint32_t materialUBOSize;              // 材质 UBO 大小
  uint32_t materialUBOOffset;            // 材质 UBO 在全局缓冲区中的偏移
  BufferHandle instanceBuffer;           // 实例数据缓冲区（实例化渲染使用）
  uint32_t instanceBufferStride;         // 实例数据步长
  
  // 变换和颜色数据（用于设置 shader uniform）
  Mat4 modelMatrix;                      // 模型矩阵
  Color color;                           // 颜色

  DrawCommand()
      : vertexCount(0), indexCount(0), instanceCount(1), textureCount(0),
        materialUBOSize(0), materialUBOOffset(0), instanceBufferStride(0),
        modelMatrix(glm::identity<Mat4>()), color(Color::White) {}

  // 检查是否使用索引绘制
  bool isIndexed() const { return indexCount > 0; }

  // 检查是否实例化绘制
  bool isInstanced() const { return instanceCount > 1; }
};

/**
 * @brief 命令批次
 *
 * 合并具有相同管线和纹理的绘制命令
 */
struct CommandBatch {
  PipelineHandle pipeline;               // 共享管线
  std::array<TextureHandle, 8> textures; // 共享纹理
  uint32_t textureCount;
  uint32_t startIndex; // 批次起始命令索引
  uint32_t count;      // 批次命令数量

  CommandBatch() : textureCount(0), startIndex(0), count(0) {}

  // 检查是否与命令兼容
  bool isCompatibleWith(const DrawCommand &cmd) const {
    if (pipeline != cmd.pipeline)
      return false;
    if (textureCount != cmd.textureCount)
      return false;
    for (uint32_t i = 0; i < textureCount; ++i) {
      if (textures[i] != cmd.textures[i])
        return false;
    }
    return true;
  }
};

/**
 * @brief 命令排序器
 *
 * 对绘制命令进行排序以优化渲染性能
 * 排序优先级：材质（管线）> 深度 > 层
 * 使用对象池复用命令内存，减少每帧分配
 */
class CommandSorter {
public:
  /**
   * @brief 默认构造函数
   */
  CommandSorter();

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

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

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

  /**
   * @brief 添加绘制命令
   * @param cmd 绘制命令
   * @return 命令索引
   */
  uint32_t addCommand(const DrawCommand &cmd);

  /**
   * @brief 排序所有命令
   */
  void sort();

  /**
   * @brief 获取排序后的命令
   * @param index 命令索引
   * @return 命令引用
   */
  const DrawCommand &getCommand(uint32_t index) const {
    return commands_[sortedIndices_[index]];
  }

  /**
   * @brief 获取命令数量
   * @return 命令数量
   */
  uint32_t getCount() const { return commandCount_; }

  /**
   * @brief 清空所有命令（保留内存）
   */
  void clear();

  /**
   * @brief 预分配命令缓冲区
   * @param capacity 预分配的容量
   */
  void reserve(uint32_t capacity);

private:
  std::vector<DrawCommand> commands_;       // 命令缓冲区（复用）
  std::vector<uint32_t> sortedIndices_;     // 排序索引缓冲区（复用）
  uint32_t commandCount_ = 0;               // 当前命令数量
  static constexpr uint32_t INITIAL_CAPACITY = 1024; // 初始容量
};

/**
 * @brief 命令批处理器
 *
 * 将连续的兼容命令合并为批次
 * 减少状态切换开销
 */
class CommandBatcher {
public:
  /**
   * @brief 处理命令列表，生成批次
   * @param sorter 已排序的命令排序器
   */
  void process(const CommandSorter &sorter);

  /**
   * @brief 获取批次数量
   * @return 批次数量
   */
  uint32_t getBatchCount() const {
    return static_cast<uint32_t>(batches_.size());
  }

  /**
   * @brief 获取批次
   * @param index 批次索引
   * @return 批次引用
   */
  const CommandBatch &getBatch(uint32_t index) const { return batches_[index]; }

  /**
   * @brief 获取批次中的命令
   * @param batchIndex 批次索引
   * @param cmdIndex 命令在批次中的索引
   * @return 命令引用
   */
  const DrawCommand &getCommand(uint32_t batchIndex, uint32_t cmdIndex) const;

  /**
   * @brief 清空所有批次
   */
  void clear();

private:
  std::vector<CommandBatch> batches_;
  const CommandSorter *sorter_ = nullptr;
};

/**
 * @brief 命令队列
 *
 * 管理一帧的所有渲染命令
 * 提供提交、排序、批处理和执行功能
 */
class CommandQueue {
public:
  /**
   * @brief 默认构造函数
   */
  CommandQueue();

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

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

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

  /**
   * @brief 初始化命令队列
   * @return 初始化是否成功
   */
  bool initialize();

  /**
   * @brief 关闭命令队列
   */
  void shutdown();

  /**
   * @brief 开始一帧的录制
   */
  void beginFrame();

  /**
   * @brief 结束一帧的录制
   */
  void endFrame();

  /**
   * @brief 提交绘制命令
   * @param material 材质
   * @param mesh 网格
   * @param transform 变换
   * @param color 颜色
   */
  void submitDraw(Ptr<Material> material, Ptr<Mesh> mesh,
                  const struct Transform &transform, const Color &color);

  /**
   * @brief 提交实例化绘制命令
   * @param material 材质
   * @param mesh 网格
   * @param instanceBuffer 实例数据缓冲区
   * @param instanceCount 实例数量
   */
  void submitDrawInstanced(Ptr<Material> material, Ptr<Mesh> mesh,
                           BufferHandle instanceBuffer, uint32_t instanceCount);

  /**
   * @brief 提交清除命令
   * @param color 清除颜色
   * @param flags 清除标志
   */
  void submitClear(const Color &color, uint32_t flags);

  /**
   * @brief 设置视口
   * @param x 视口 X
   * @param y 视口 Y
   * @param width 视口宽度
   * @param height 视口高度
   */
  void setViewport(int32_t x, int32_t y, int32_t width, int32_t height);

  /**
   * @brief 执行所有命令
   *
   * 排序、批处理并执行所有提交的命令
   * @param frameIndex 当前帧索引（用于双缓冲）
   */
  void execute(uint32_t frameIndex);

  /**
   * @brief 更新全局 UBO 数据
   * @param viewProjection 视图投影矩阵
   * @param deltaTime 帧时间
   * @param screenWidth 屏幕宽度
   * @param screenHeight 屏幕高度
   * @param frameIndex 当前帧索引（用于双缓冲）
   */
  void updateGlobalUBO(const Mat4& viewProjection, float deltaTime,
                       uint32_t screenWidth, uint32_t screenHeight, uint32_t frameIndex);

  /**
   * @brief 获取当前命令数量
   * @return 命令数量
   */
  uint32_t getCommandCount() const { return sorter_.getCount(); }

  /**
   * @brief 获取命令列表
   * @return 命令列表指针
   */
  RHICommandList* getCommandList() const { return commandList_.get(); }

  /**
   * @brief 获取渲染统计信息
   * @return 渲染统计
   */
  const RenderStats& getStats() const { return stats_; }

  /**
   * @brief 重置统计信息
   */
  void resetStats() { stats_ = {}; }

private:
  CommandSorter sorter_;
  CommandBatcher batcher_;
  RHIContext *context_ = nullptr;
  std::unique_ptr<RHICommandList> commandList_;

  // UBO 管理器
  std::unique_ptr<UniformBufferManager> uboManager_;

  // 渲染统计
  RenderStats stats_;

  // 全局 UBO 数据 - 必须与着色器中的 std140 布局完全匹配
  // layout(std140, binding = 0) uniform GlobalUBO {
  //     mat4 uViewProjection;    // 64 bytes, offset 0
  //     vec4 uCameraPosition;    // 16 bytes, offset 64
  //     float uTime;             // 4 bytes, offset 80
  //     float uDeltaTime;        // 4 bytes, offset 84
  //     vec2 uScreenSize;        // 8 bytes, offset 88
  // }; // 总大小: 96 bytes (std140 对齐)
  struct alignas(16) GlobalUBOData {
    float viewProjection[16];  // 64 bytes, offset 0
    float cameraPosition[4];   // 16 bytes, offset 64
    float time;                // 4 bytes, offset 80
    float deltaTime;           // 4 bytes, offset 84
    float screenSize[2];       // 8 bytes, offset 88
  } globalUBOData_;

  // 材质 UBO 数据缓冲区
  std::vector<uint8_t> materialUBOData_;

  // 当前材质 ID 计数器
  uint32_t nextMaterialId_ = 1;

  // 材质到 ID 的映射
  std::unordered_map<Material *, uint32_t> materialIds_;

  // 当前材质 UBO 缓冲区
  UniformBuffer* currentMaterialUBO_ = nullptr;
  uint32_t materialUBOBufferOffset_ = 0;

  /**
   * @brief 获取或创建材质 ID
   * @param material 材质指针
   * @return 材质 ID
   */
  uint32_t getMaterialId(Material *material);

  /**
   * @brief 执行单个批次
   * @param batchIndex 批次索引
   * @param batch 命令批次
   * @param frameIndex 当前帧索引（用于双缓冲）
   */
  void executeBatch(uint32_t batchIndex, const CommandBatch &batch, uint32_t frameIndex);

  /**
   * @brief 分配材质 UBO 空间（使用 CPU 缓冲区）
   * @param size 需要的空间大小
   * @return 分配的 UBO 指针（nullptr 表示使用 CPU 缓冲区）和偏移量
   */
  std::pair<UniformBuffer*, uint32_t> allocateMaterialUBO(uint32_t size);

  /**
   * @brief 将材质 UBO CPU 缓冲区数据刷新到 GPU
   */
  void flushMaterialUBOToGPU();

  /**
   * @brief 获取当前材质 UBO
   * @return 当前材质 UBO 指针
   */
  UniformBuffer* getCurrentMaterialUBO() const;
};

} // namespace extra2d