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 <types/math/color.h>
#include <types/math/transform.h>
#include <unordered_map>
#include <vector>

namespace extra2d {

// 前向声明
class CommandQueue;
class Material;
class Mesh;
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 大小
  
  // 变换和颜色数据（用于设置 shader uniform）
  Mat4 modelMatrix;                      // 模型矩阵
  Color color;                           // 颜色

  DrawCommand()
      : vertexCount(0), indexCount(0), instanceCount(1), textureCount(0),
        materialUBOSize(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 添加绘制命令
   * @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 static_cast<uint32_t>(commands_.size()); }

  /**
   * @brief 清空所有命令
   */
  void clear();

private:
  std::vector<DrawCommand> commands_;
  std::vector<uint32_t> sortedIndices_;
};

/**
 * @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 instanceCount 实例数量
   */
  void submitDrawInstanced(Ptr<Material> material, Ptr<Mesh> mesh,
                           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 执行所有命令
   *
   * 排序、批处理并执行所有提交的命令
   */
  void execute();

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

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

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

  // 全局 UBO 数据
  struct GlobalUBOData {
    float viewProjection[16];
    float time;
    float screenSize[2];
    float padding;
  } globalUBOData_;

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

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

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

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

  /**
   * @brief 执行单个批次
   * @param batchIndex 批次索引
   * @param batch 命令批次
   */
  void executeBatch(uint32_t batchIndex, const CommandBatch &batch);
};

} // namespace extra2d