#include "renderer/scene_renderer.hpp"

#include <algorithm>
#include <sophia/components.hpp>
#include <sophia/util/assert.hpp>
#include <tracy/Tracy.hpp>
#include <tuple>

namespace sophia {

SceneRenderer::SceneRenderer(const SceneRegistry& sceneRegistry,
                             ResourceRegistry& resources, Device& device,
                             vk::Format colorFormat)
    : sceneRegistry{sceneRegistry},
      resources{resources},
      spriteRenderSystem{device, colorFormat, resources} {}

void SceneRenderer::beginFrame(u32 frameIndex) {
  ZoneScoped;
  this->spriteRenderSystem.beginFrame(frameIndex);
}

void SceneRenderer::render(vk::CommandBuffer commandBuffer,
                           const SceneView& sceneView, Entity camera) {
  ZoneScoped;
  SOPHIA_ASSERT(camera.hasComponent<CameraComponent>(),
                "camera does not have camera component!");

  vk::Extent2D extent =
      this->resources.renderTargetPool.get(sceneView.target).getExtent();
  float aspect =
      static_cast<float>(extent.width) / static_cast<float>(extent.height);

  glm::mat4 viewProjection =
      camera.getComponent<CameraComponent>().getViewProjection(
          camera.getComponent<TransformComponent>(), aspect);

  renderSprites(commandBuffer, this->sceneRegistry.get(sceneView.scene),
                viewProjection);

  // TODO: meshes. Group MeshComponent entities by mesh and draw each group
  // (instanced), as sprites are grouped by texture.
}

void SceneRenderer::renderSprites(vk::CommandBuffer commandBuffer,
                                  const Scene& scene,
                                  const glm::mat4& viewProjection) {
  this->spriteDrawItems.clear();

  auto view =
      scene.registry.view<const TransformComponent, const SpriteComponent>();
  for (auto [entity, transform, sprite] : view.each()) {
    TextureHandle texture = this->resources.texturePool.isAlive(sprite.texture)
                                ? sprite.texture
                                : this->resources.whiteTexture;

    this->spriteDrawItems.push_back(
        {sprite.sortOrder,
         transform.position.z,
         texture,
         {sprite.getModel(transform), sprite.color, sprite.uvRect}});
  }

  // Texture is the last key so that sprites whose order doesn't matter
  // relative to each other still end up in as few batches as possible.
  std::ranges::sort(this->spriteDrawItems, {}, [](const SpriteDrawItem& item) {
    return std::tuple{item.sortOrder, item.depth, item.texture};
  });

  this->spriteInstances.clear();
  this->spriteBatches.clear();
  for (const SpriteDrawItem& item : this->spriteDrawItems) {
    auto index = static_cast<u32>(this->spriteInstances.size());
    this->spriteInstances.push_back(item.instance);

    if (this->spriteBatches.empty() ||
        this->spriteBatches.back().texture != item.texture) {
      this->spriteBatches.push_back({item.texture, index, 0});
    }
    this->spriteBatches.back().instanceCount++;
  }

  this->spriteRenderSystem.render(commandBuffer, viewProjection,
                                  this->spriteInstances, this->spriteBatches);
}

}  // namespace sophia