#include "renderer/sprite_render_system.hpp"

#include <sophia/util/assert.hpp>
#include <sophia/util/logger.hpp>
#include <tracy/Tracy.hpp>

#include "shaders/sprite.frag.spv.h"
#include "shaders/sprite.vert.spv.h"
#include "vulkan/descriptors/descriptor_writer.hpp"

namespace sophia {

namespace {
constexpr u32 MAX_TEXTURES = 1024;
constexpr u32 VERTICES_PER_SPRITE = 6;
}  // namespace

SpriteRenderSystem::SpriteRenderSystem(Device& device, vk::Format colorFormat,
                                       ResourceRegistry& resources)
    : device{device}, resources{resources} {
  createDescriptors();
  createPipelineLayout();
  createPipeline(colorFormat);
  createInstanceBuffers();
}

SpriteRenderSystem::~SpriteRenderSystem() {
  log::trace("destroyed sophia::SpriteRenderSystem");
}

void SpriteRenderSystem::beginFrame(u32 frameIndex) {
  this->frameIndex = frameIndex;
  this->instanceCursor = 0;
}

void SpriteRenderSystem::render(vk::CommandBuffer commandBuffer,
                                const glm::mat4& viewProjection,
                                std::span<const Instance> instances,
                                std::span<const Batch> batches) {
  ZoneScoped;
  if (instances.empty()) {
    return;
  }

  u32 available = MAX_INSTANCES_PER_FRAME - this->instanceCursor;
  if (instances.size() > available) {
    log::warning("SpriteRenderSystem: frame instance buffer full, dropping",
                 instances.size() - available, "sprites");
  }

  Buffer& buffer = *this->instanceBuffers[this->frameIndex];
  u32 base = this->instanceCursor;
  u32 count = std::min(static_cast<u32>(instances.size()), available);
  buffer.writeToBuffer(const_cast<Instance*>(instances.data()),
                       count * sizeof(Instance), base * sizeof(Instance));
  this->instanceCursor += count;

  this->graphicsPipeline->bind(commandBuffer);

  PushConstantData push{viewProjection};
  commandBuffer.pushConstants(this->pipelineLayout.get(),
                              vk::ShaderStageFlagBits::eVertex, 0,
                              sizeof(PushConstantData), &push);

  vk::Buffer vertexBuffer = buffer.getBuffer();
  vk::DeviceSize offset = 0;
  commandBuffer.bindVertexBuffers(0, 1, &vertexBuffer, &offset);

  for (const Batch& batch : batches) {
    if (batch.firstInstance >= count) {
      break;
    }
    u32 batchCount = std::min(batch.instanceCount, count - batch.firstInstance);

    vk::DescriptorSet set = getDescriptorSet(batch.texture);
    commandBuffer.bindDescriptorSets(vk::PipelineBindPoint::eGraphics,
                                     this->pipelineLayout.get(), 0, 1, &set, 0,
                                     nullptr);
    commandBuffer.draw(VERTICES_PER_SPRITE, batchCount, 0,
                       base + batch.firstInstance);
  }
}

vk::DescriptorSet SpriteRenderSystem::getDescriptorSet(TextureHandle handle) {
  auto it = this->descriptorSets.find(handle.index);
  if (it != this->descriptorSets.end() &&
      it->second.generation == handle.generation) {
    return it->second.set;
  }

  const Texture& texture = this->resources.texturePool.get(handle);
  vk::DescriptorImageInfo imageInfo = texture.descriptorInfo();
  DescriptorWriter writer{*this->setLayout, *this->descriptorPool};
  writer.writeImage(0, &imageInfo);

  if (it == this->descriptorSets.end()) {
    vk::DescriptorSet set;
    bool allocated = writer.build(set);
    SOPHIA_ASSERT(allocated, "SpriteRenderSystem: out of texture descriptors");
    this->descriptorSets.emplace(static_cast<u32>(handle.index),
                                 CachedDescriptorSet{handle.generation, set});
    return set;
  }

  // The slot was reused by a new texture. Releasing a texture waits for the
  // device to go idle, so no in-flight frame can still be using this set.
  writer.overwrite(it->second.set);
  it->second.generation = handle.generation;
  return it->second.set;
}

void SpriteRenderSystem::createDescriptors() {
  this->setLayout =
      DescriptorSetLayout::Builder(this->device)
          .addBinding(0, vk::DescriptorType::eCombinedImageSampler,
                      vk::ShaderStageFlagBits::eFragment)
          .build();

  this->descriptorPool =
      DescriptorPool::Builder(this->device)
          .setMaxSets(MAX_TEXTURES)
          .addPoolSize(vk::DescriptorType::eCombinedImageSampler, MAX_TEXTURES)
          .build();
}

void SpriteRenderSystem::createPipelineLayout() {
  vk::PushConstantRange pushConstantRange{};
  pushConstantRange.stageFlags = vk::ShaderStageFlagBits::eVertex;
  pushConstantRange.offset = 0;
  pushConstantRange.size = sizeof(PushConstantData);

  vk::DescriptorSetLayout setLayout = this->setLayout->getDescriptorSetLayout();

  vk::PipelineLayoutCreateInfo pipelineLayoutInfo = {};
  pipelineLayoutInfo.setLayoutCount = 1;
  pipelineLayoutInfo.pSetLayouts = &setLayout;
  pipelineLayoutInfo.pushConstantRangeCount = 1;
  pipelineLayoutInfo.pPushConstantRanges = &pushConstantRange;

  try {
    this->pipelineLayout =
        this->device.get()->createPipelineLayoutUnique(pipelineLayoutInfo);
    log::trace("created vk::PipelineLayout");
  } catch (const vk::SystemError& err) {
    log::fatal("failed to create vk::PipelineLayout");
    throw std::runtime_error("failed to create vk::PipelineLayout");
  }
}

void SpriteRenderSystem::createPipeline(vk::Format colorFormat) {
  auto vertexShader = Shader::Builder()
                          .fromSPIRV(shaders::sprite_vert_spv)
                          .setStage(ShaderStage::VERTEX)
                          .build();
  auto fragmentShader = Shader::Builder()
                            .fromSPIRV(shaders::sprite_frag_spv)
                            .setStage(ShaderStage::FRAGMENT)
                            .build();

  std::vector<vk::VertexInputBindingDescription> bindings{
      {0, sizeof(Instance), vk::VertexInputRate::eInstance}};

  std::vector<vk::VertexInputAttributeDescription> attributes;
  // A mat4 attribute takes four consecutive locations, one per column.
  for (u32 column = 0; column < 4; column++) {
    attributes.push_back({column, 0, vk::Format::eR32G32B32A32Sfloat,
                          static_cast<u32>(offsetof(Instance, model) +
                                           column * sizeof(glm::vec4))});
  }
  attributes.push_back({4, 0, vk::Format::eR32G32B32A32Sfloat,
                        static_cast<u32>(offsetof(Instance, color))});
  attributes.push_back({5, 0, vk::Format::eR32G32B32A32Sfloat,
                        static_cast<u32>(offsetof(Instance, uvRect))});

  this->graphicsPipeline = GraphicsPipeline::Builder(this->device)
                               .setVertexShader(*vertexShader)
                               .setFragmentShader(*fragmentShader)
                               .setPipelineLayout(this->pipelineLayout.get())
                               .setColorAttachmentFormats({colorFormat})
                               .setVertexInput(bindings, attributes)
                               .setAlphaBlending()
                               .build();
}

void SpriteRenderSystem::createInstanceBuffers() {
  for (std::unique_ptr<Buffer>& buffer : this->instanceBuffers) {
    buffer = std::make_unique<Buffer>(
        this->device, sizeof(Instance), MAX_INSTANCES_PER_FRAME,
        vk::BufferUsageFlagBits::eVertexBuffer,
        vk::MemoryPropertyFlagBits::eHostVisible |
            vk::MemoryPropertyFlagBits::eHostCoherent);
    buffer->map();
  }
}

}  // namespace sophia