Extra2D/src/renderer/command_queue.cpp

#include <algorithm>
#include <cstring>
#include <vector>
#include <renderer/command_queue.h>
#include <renderer/instance_buffer.h>
#include <renderer/material.h>
#include <renderer/mesh.h>
#include <renderer/rhi_module.h>
#include <renderer/rhi/opengl/gl_command_list.h>
#include <renderer/rhi/opengl/gl_texture.h>
#include <types/math/transform.h>
#include <types/ptr/intrusive_ptr.h>
#include <utils/logger.h>

#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/type_ptr.hpp>

namespace extra2d {

// ========================================
// CommandSorter 实现
// ========================================

uint32_t CommandSorter::addCommand(const DrawCommand &cmd) {
  uint32_t index = static_cast<uint32_t>(commands_.size());
  commands_.push_back(cmd);
  sortedIndices_.push_back(index);
  return index;
}

void CommandSorter::sort() {
  // 使用稳定排序保持相同键的命令顺序
  std::stable_sort(sortedIndices_.begin(), sortedIndices_.end(),
                   [this](uint32_t a, uint32_t b) {
                     return commands_[a].key < commands_[b].key;
                   });
}

void CommandSorter::clear() {
  commands_.clear();
  sortedIndices_.clear();
}

// ========================================
// CommandBatcher 实现
// ========================================

void CommandBatcher::process(const CommandSorter &sorter) {
  clear();
  sorter_ = &sorter;

  if (sorter.getCount() == 0) {
    return;
  }

  CommandBatch currentBatch;
  currentBatch.startIndex = 0;
  currentBatch.count = 1;

  const auto &firstCmd = sorter.getCommand(0);
  currentBatch.pipeline = firstCmd.pipeline;
  currentBatch.textureCount = firstCmd.textureCount;
  for (uint32_t i = 0; i < firstCmd.textureCount; ++i) {
    currentBatch.textures[i] = firstCmd.textures[i];
  }

  // 遍历所有命令，合并兼容的命令
  for (uint32_t i = 1; i < sorter.getCount(); ++i) {
    const auto &cmd = sorter.getCommand(i);

    if (currentBatch.isCompatibleWith(cmd)) {
      // 兼容，加入当前批次
      currentBatch.count++;
    } else {
      // 不兼容，保存当前批次并创建新批次
      batches_.push_back(currentBatch);

      currentBatch.startIndex = i;
      currentBatch.count = 1;
      currentBatch.pipeline = cmd.pipeline;
      currentBatch.textureCount = cmd.textureCount;
      for (uint32_t j = 0; j < cmd.textureCount; ++j) {
        currentBatch.textures[j] = cmd.textures[j];
      }
    }
  }

  // 保存最后一个批次
  batches_.push_back(currentBatch);
}

const DrawCommand &CommandBatcher::getCommand(uint32_t batchIndex,
                                              uint32_t cmdIndex) const {
  const auto &batch = batches_[batchIndex];
  return sorter_->getCommand(batch.startIndex + cmdIndex);
}

void CommandBatcher::clear() {
  batches_.clear();
  sorter_ = nullptr;
}

// ========================================
// CommandQueue 实现
// ========================================

CommandQueue::CommandQueue() = default;

CommandQueue::~CommandQueue() { shutdown(); }

CommandQueue::CommandQueue(CommandQueue &&other) noexcept
    : sorter_(std::move(other.sorter_)), batcher_(std::move(other.batcher_)),
      context_(other.context_), commandList_(std::move(other.commandList_)),
      uboManager_(std::move(other.uboManager_)),
      globalUBOData_(other.globalUBOData_),
      materialUBOData_(std::move(other.materialUBOData_)),
      nextMaterialId_(other.nextMaterialId_),
      materialIds_(std::move(other.materialIds_)),
      currentMaterialUBO_(other.currentMaterialUBO_),
      currentMaterialUBOOffset_(other.currentMaterialUBOOffset_) {
  other.context_ = nullptr;
  other.globalUBOData_ = {};
  other.nextMaterialId_ = 1;
  other.currentMaterialUBO_ = nullptr;
  other.currentMaterialUBOOffset_ = 0;
}

CommandQueue &CommandQueue::operator=(CommandQueue &&other) noexcept {
  if (this != &other) {
    shutdown();

    sorter_ = std::move(other.sorter_);
    batcher_ = std::move(other.batcher_);
    context_ = other.context_;
    commandList_ = std::move(other.commandList_);
    uboManager_ = std::move(other.uboManager_);
    globalUBOData_ = other.globalUBOData_;
    materialUBOData_ = std::move(other.materialUBOData_);
    nextMaterialId_ = other.nextMaterialId_;
    materialIds_ = std::move(other.materialIds_);
    currentMaterialUBO_ = other.currentMaterialUBO_;
    currentMaterialUBOOffset_ = other.currentMaterialUBOOffset_;

    other.context_ = nullptr;
    other.globalUBOData_ = {};
    other.nextMaterialId_ = 1;
    other.currentMaterialUBO_ = nullptr;
    other.currentMaterialUBOOffset_ = 0;
  }
  return *this;
}

bool CommandQueue::initialize() {
  // 获取 RHI 模块
  auto *rhiModule = RHIModule::get();
  if (!rhiModule) {
    E2D_LOG_ERROR("RHIModule not available");
    return false;
  }

  auto *device = rhiModule->getDevice();
  if (!device) {
    E2D_LOG_ERROR("RHIDevice not available");
    return false;
  }

  context_ = rhiModule->getContext();
  if (!context_) {
    E2D_LOG_ERROR("RHIContext not available");
    return false;
  }

  // 创建命令列表
  commandList_ = device->createCommandList();
  if (!commandList_) {
    E2D_LOG_ERROR("Failed to create command list");
    return false;
  }

  // 初始化 UBO 管理器
  uboManager_ = std::make_unique<UniformBufferManager>();
  if (!uboManager_->initialize()) {
    E2D_LOG_ERROR("Failed to initialize UniformBufferManager");
    return false;
  }

  // 预分配材质 UBO 数据缓冲区
  materialUBOData_.reserve(1024 * 1024); // 1MB

  E2D_LOG_INFO("CommandQueue initialized");
  return true;
}

void CommandQueue::shutdown() {
  uboManager_.reset();
  commandList_.reset();
  context_ = nullptr;
  materialUBOData_.clear();
  E2D_LOG_INFO("CommandQueue shutdown");
}

void CommandQueue::beginFrame() {
  sorter_.clear();
  batcher_.clear();
  materialIds_.clear();
  nextMaterialId_ = 1;
  materialUBOData_.clear();

  // 重置材质 UBO 分配状态
  currentMaterialUBO_ = nullptr;
  currentMaterialUBOOffset_ = 0;

  // 重置 UBO 管理器的材质 UBO 池
  if (uboManager_) {
    uboManager_->resetMaterialUBOs();
  }

  // 开始录制命令
  if (commandList_) {
    commandList_->begin();
  }
}

void CommandQueue::endFrame() {
  // 结束录制
  if (commandList_) {
    commandList_->end();
  }
}

uint32_t CommandQueue::getMaterialId(Material *material) {
  auto it = materialIds_.find(material);
  if (it != materialIds_.end()) {
    return it->second;
  }

  uint32_t id = nextMaterialId_++;
  materialIds_[material] = id;
  return id;
}

std::pair<UniformBuffer*, uint32_t> CommandQueue::allocateMaterialUBO(uint32_t size) {
  if (!uboManager_ || size == 0) {
    return {nullptr, 0};
  }

  // 如果当前 UBO 没有足够的空间，获取一个新的
  if (currentMaterialUBO_ == nullptr || 
      currentMaterialUBOOffset_ + size > currentMaterialUBO_->getSize()) {
    currentMaterialUBO_ = uboManager_->acquireMaterialUBO(size);
    currentMaterialUBOOffset_ = 0;
  }

  if (!currentMaterialUBO_) {
    return {nullptr, 0};
  }

  uint32_t offset = currentMaterialUBOOffset_;
  currentMaterialUBOOffset_ += size;

  // 对齐到 16 字节（std140 要求）
  currentMaterialUBOOffset_ = (currentMaterialUBOOffset_ + 15) & ~15;

  return {currentMaterialUBO_, offset};
}

void CommandQueue::submitDraw(Ptr<Material> material, Ptr<Mesh> mesh,
                              const struct Transform &transform,
                              const Color &color) {
  if (!material || !mesh || !material->getShader()) {
    return;
  }

  DrawCommand cmd;

  // 构建排序键
  uint32_t materialId = getMaterialId(material.get());
  cmd.key = DrawKey::make(materialId, 0, 0);

  // 设置管线
  cmd.pipeline = material->getPipeline();

  // 设置网格数据
  cmd.vertexBuffer = mesh->getVertexBuffer();
  cmd.indexBuffer = mesh->getIndexBuffer();
  cmd.vertexCount = mesh->getVertexCount();
  cmd.indexCount = mesh->getIndexCount();

  // 设置变换和颜色
  float matrixData[16];
  transform.toMatrix(matrixData);
  cmd.modelMatrix = glm::make_mat4(matrixData);
  cmd.color = color;

  // 设置纹理
  const auto &textures = material->getTextures();
  cmd.textureCount =
      static_cast<uint32_t>(std::min(textures.size(), size_t(8)));
  for (uint32_t i = 0; i < cmd.textureCount; ++i) {
    if (textures[i].texture) {
      cmd.textures[i] = textures[i].texture->getHandle();
    }
  }

  // 分配材质 UBO 空间并更新数据
  uint32_t materialDataSize = material->getDataSize();
  if (materialDataSize > 0) {
    auto [ubo, offset] = allocateMaterialUBO(materialDataSize);
    if (ubo) {
      // 复制材质数据到临时缓冲区，以便修改颜色
      std::vector<uint8_t> uboData(materialDataSize);
      std::memcpy(uboData.data(), material->getData(), materialDataSize);

      // 将实例颜色应用到 UBO 数据中的 uColor 参数
      auto layout = material->getLayout();
      if (layout) {
        const auto* param = layout->getParam("uColor");
        if (param && param->type == MaterialParamType::Color) {
          std::memcpy(uboData.data() + param->offset, &color.r, sizeof(float) * 4);
        }
      }

      ubo->update(uboData.data(), materialDataSize, offset);
      cmd.materialUBO = BufferHandle(ubo->getRHIBuffer());
      cmd.materialUBOSize = materialDataSize;
      cmd.materialUBOOffset = offset;
    }
  }

  sorter_.addCommand(cmd);
}

void CommandQueue::submitDrawInstanced(Ptr<Material> material, Ptr<Mesh> mesh,
                                       BufferHandle instanceBuffer, uint32_t instanceCount) {
  if (!material || !mesh || !material->getShader() || instanceCount == 0) {
    return;
  }

  DrawCommand cmd;

  // 构建排序键
  uint32_t materialId = getMaterialId(material.get());
  cmd.key = DrawKey::make(materialId, 0, 0);

  // 设置管线
  cmd.pipeline = material->getPipeline();

  // 设置网格数据
  cmd.vertexBuffer = mesh->getVertexBuffer();
  cmd.indexBuffer = mesh->getIndexBuffer();
  cmd.vertexCount = mesh->getVertexCount();
  cmd.indexCount = mesh->getIndexCount();
  cmd.instanceCount = instanceCount;

  // 设置实例缓冲区
  cmd.instanceBuffer = instanceBuffer;
  cmd.instanceBufferStride = sizeof(InstanceData);  // 64 bytes

  // 设置纹理
  const auto &textures = material->getTextures();
  cmd.textureCount =
      static_cast<uint32_t>(std::min(textures.size(), size_t(8)));
  for (uint32_t i = 0; i < cmd.textureCount; ++i) {
    if (textures[i].texture) {
      cmd.textures[i] = textures[i].texture->getHandle();
    }
  }

  // 分配材质 UBO 空间并更新数据
  uint32_t materialDataSize = material->getDataSize();
  if (materialDataSize > 0) {
    auto [ubo, offset] = allocateMaterialUBO(materialDataSize);
    if (ubo) {
      ubo->update(material->getData(), materialDataSize, offset);
      cmd.materialUBO = BufferHandle(ubo->getRHIBuffer());
      cmd.materialUBOSize = materialDataSize;
      cmd.materialUBOOffset = offset;
    }
  }

  sorter_.addCommand(cmd);
}

void CommandQueue::submitClear(const Color &color, uint32_t flags) {
  // 清除命令直接执行，不加入排序队列
  if (!commandList_) {
    return;
  }

  commandList_->clear(static_cast<ClearFlags>(flags), color, 1.0f, 0);
}

void CommandQueue::setViewport(int32_t x, int32_t y, int32_t width,
                               int32_t height) {
  if (!commandList_) {
    return;
  }

  Viewport viewport;
  viewport.x = static_cast<float>(x);
  viewport.y = static_cast<float>(y);
  viewport.width = static_cast<float>(width);
  viewport.height = static_cast<float>(height);
  viewport.minDepth = 0.0f;
  viewport.maxDepth = 1.0f;

  commandList_->setViewport(viewport);
}

void CommandQueue::updateGlobalUBO(const Mat4& viewProjection, float deltaTime,
                                   uint32_t screenWidth, uint32_t screenHeight) {
  if (!uboManager_) {
    E2D_LOG_WARN("CommandQueue::updateGlobalUBO: uboManager is null");
    return;
  }

  // 填充全局 UBO 数据
  std::memcpy(globalUBOData_.viewProjection, glm::value_ptr(viewProjection), sizeof(float) * 16);
  globalUBOData_.cameraPosition[0] = 0.0f;
  globalUBOData_.cameraPosition[1] = 0.0f;
  globalUBOData_.cameraPosition[2] = 0.0f;
  globalUBOData_.cameraPosition[3] = 1.0f;
  globalUBOData_.time = 0.0f;  // TODO: 传递实际时间
  globalUBOData_.deltaTime = deltaTime;
  globalUBOData_.screenSize[0] = static_cast<float>(screenWidth);
  globalUBOData_.screenSize[1] = static_cast<float>(screenHeight);

  // 更新 UBO
  uboManager_->updateGlobalUBO(&globalUBOData_, sizeof(globalUBOData_));
}

void CommandQueue::execute() {
  if (!commandList_) {
    E2D_LOG_ERROR("CommandQueue::execute: commandList is null");
    return;
  }

  // 排序命令
  sorter_.sort();

  // 批处理
  batcher_.process(sorter_);

  // 执行批次
  for (uint32_t i = 0; i < batcher_.getBatchCount(); ++i) {
    executeBatch(i, batcher_.getBatch(i));
  }

  // 提交命令到 GPU
  commandList_->submit();
}

void CommandQueue::executeBatch(uint32_t batchIndex, const CommandBatch &batch) {
  if (!commandList_) {
    return;
  }

  // 绑定管线
  if (batch.pipeline.isValid()) {
    commandList_->setPipeline(batch.pipeline.get());
  } else {
    E2D_LOG_WARN("Batch has no valid pipeline!");
  }

  // 绑定全局 UBO (binding = 0)
  if (uboManager_) {
    UniformBuffer* globalUBO = uboManager_->getGlobalUBO();
    if (globalUBO) {
      commandList_->setUniformBuffer(0, globalUBO->getRHIBuffer());
    }
  }

  // 绑定纹理
  if (batch.textureCount > 0) {
    for (uint32_t i = 0; i < batch.textureCount; ++i) {
      if (batch.textures[i].isValid()) {
        commandList_->setTexture(i, batch.textures[i].get());
      }
    }
  } else {
    // 如果没有纹理，绑定默认的白色纹理
    auto* rhiModule = RHIModule::get();
    if (rhiModule && rhiModule->getDevice()) {
      // 使用默认纹理（白色像素）
      // 注意：这里应该使用 AssetsModule 的默认纹理
      // 暂时跳过，因为我们需要访问 AssetsModule
    }
  }

  // 执行批次中的所有命令
  for (uint32_t i = 0; i < batch.count; ++i) {
    const auto &cmd = batcher_.getCommand(batchIndex, i);

    // 绑定顶点缓冲区
    if (cmd.vertexBuffer.isValid()) {
      commandList_->setVertexBuffer(0, cmd.vertexBuffer.get(), 0);
    } else {
      E2D_LOG_WARN("Draw command has no valid vertex buffer!");
    }

    // 绑定实例缓冲区（实例化渲染）
    if (cmd.isInstanced() && cmd.instanceBuffer.isValid()) {
      commandList_->setVertexBuffer(1, cmd.instanceBuffer.get(), 0, cmd.instanceBufferStride);
    }

    // 绑定索引缓冲区（如果有）
    if (cmd.isIndexed() && cmd.indexBuffer.isValid()) {
      commandList_->setIndexBuffer(cmd.indexBuffer.get(), IndexType::UInt16, 0);
    }

    // 绑定材质 UBO (binding = 1)
    if (cmd.materialUBO.isValid()) {
      commandList_->setUniformBuffer(1, cmd.materialUBO.get(), cmd.materialUBOOffset, cmd.materialUBOSize);
    }

    // 设置模型矩阵（仅在非实例化渲染时使用）
    if (!cmd.isInstanced()) {
      commandList_->setUniform("uModelMatrix", cmd.modelMatrix);
    }

    // 绘制
    if (cmd.isInstanced()) {
      if (cmd.isIndexed()) {
        commandList_->drawIndexed(cmd.indexCount, 0, 0, cmd.instanceCount, 0);
      } else {
        commandList_->draw(cmd.vertexCount, 0, cmd.instanceCount, 0);
      }
    } else {
      if (cmd.isIndexed()) {
        commandList_->drawIndexed(cmd.indexCount, 0, 0, 1, 0);
      } else {
        commandList_->draw(cmd.vertexCount, 0, 1, 0);
      }
    }
  }
}

} // namespace extra2d