Extra2D/src/renderer/command_queue.cpp

#include <algorithm>
#include <cstring>
#include <renderer/command_queue.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_)),
      globalUBOData_(other.globalUBOData_),
      materialUBOData_(std::move(other.materialUBOData_)),
      nextMaterialId_(other.nextMaterialId_),
      materialIds_(std::move(other.materialIds_)) {
  other.context_ = nullptr;
  other.globalUBOData_ = {};
  other.nextMaterialId_ = 1;
}

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_);
    globalUBOData_ = other.globalUBOData_;
    materialUBOData_ = std::move(other.materialUBOData_);
    nextMaterialId_ = other.nextMaterialId_;
    materialIds_ = std::move(other.materialIds_);

    other.context_ = nullptr;
    other.globalUBOData_ = {};
    other.nextMaterialId_ = 1;
  }
  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 数据缓冲区
  materialUBOData_.reserve(1024 * 1024); // 1MB

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

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

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

  // 开始录制命令
  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;
}

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) {
    uint32_t uboOffset = static_cast<uint32_t>(materialUBOData_.size());
    materialUBOData_.resize(uboOffset + materialDataSize);
    std::memcpy(materialUBOData_.data() + uboOffset, material->getData(),
                materialDataSize);
    cmd.materialUBOSize = materialDataSize;
    // 注意：实际的 UBO 句柄需要在执行时从 UBO 管理器获取
  }

  sorter_.addCommand(cmd);
}

void CommandQueue::submitDrawInstanced(Ptr<Material> material, Ptr<Mesh> mesh,
                                       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;

  // 设置纹理
  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) {
    uint32_t uboOffset = static_cast<uint32_t>(materialUBOData_.size());
    materialUBOData_.resize(uboOffset + materialDataSize);
    std::memcpy(materialUBOData_.data() + uboOffset, material->getData(),
                materialDataSize);
    cmd.materialUBOSize = materialDataSize;
  }

  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::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!");
  }

  // 绑定纹理
  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.isIndexed() && cmd.indexBuffer.isValid()) {
      commandList_->setIndexBuffer(cmd.indexBuffer.get(), IndexType::UInt16, 0);
    }

    // 设置 shader uniform 变量
    auto* glCmdList = dynamic_cast<extra2d::GLCommandList*>(commandList_.get());
    if (glCmdList) {
      // 设置模型矩阵
      glCmdList->setUniform("uModelMatrix", cmd.modelMatrix);
      // 设置颜色
      glCmdList->setUniform("uColor", cmd.color);
      // 设置 view projection 矩阵
      Mat4 viewProj = glm::ortho(0.0f, 1280.0f, 0.0f, 720.0f, -1.0f, 1.0f);
      glCmdList->setUniform("uViewProjection", viewProj);
      // 设置色调颜色
      glCmdList->setUniform("uTintColor", Color::White);
      // 设置透明度
      glCmdList->setUniform("uOpacity", 1.0f);
    }

    // 绘制
    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