#include "model_render_system.hpp"

#include <sophia/util/assert.hpp>

#include "shaders/model_render_system.frag.spv.h"
#include "shaders/model_render_system.vert.spv.h"

namespace sophia {

ModelRenderSystem::ModelRenderSystem(Device& device,
                                     RenderContext& renderContext)
    : device{device}, renderContext{renderContext} {
  createPipelineLayout();

  SOPHIA_ASSERT(this->pipelineLayout.get() != nullptr,
                "Cannot create pipeline before pipeline layout");

  auto vertexShader = Shader::Builder()
                          .fromSPIRV(shaders::model_render_system_vert_spv)
                          .setStage(ShaderStage::VERTEX)
                          .build();
  auto fragmentShader = Shader::Builder()
                            .fromSPIRV(shaders::model_render_system_frag_spv)
                            .setStage(ShaderStage::FRAGMENT)
                            .build();

  auto bindingDescriptions = Mesh::Vertex::getBindingDescriptions();
  auto attributeDescriptions = Mesh::Vertex::getAttributeDescriptions();

  this->graphicsPipeline =
      GraphicsPipeline::Builder(device)
          .setVertexShader(*vertexShader)
          .setFragmentShader(*fragmentShader)
          .setPipelineLayout(pipelineLayout.get())
          .setColorAttachmentFormats({renderContext.getImageFormat()})
          .setVertexInput(bindingDescriptions, attributeDescriptions)
          .build();
}

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

void ModelRenderSystem::render(vk::CommandBuffer commandBuffer,
                               std::span<const MeshInstance> instances,
                               const FrameInfo& frameInfo) {
  glm::mat4 viewProjection = frameInfo.camera.getViewProjection();

  this->graphicsPipeline->bind(commandBuffer);

  for (const MeshInstance& instance : instances) {
    PushConstantData push{};
    push.transform = viewProjection * instance.transform;

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

    instance.mesh->bind(commandBuffer);
    instance.mesh->draw(commandBuffer);
  }
}

void ModelRenderSystem::createPipelineLayout() {
  vk::PushConstantRange pushConstantRange{};

  pushConstantRange.stageFlags =
      vk::ShaderStageFlagBits::eVertex | vk::ShaderStageFlagBits::eFragment;
  pushConstantRange.offset = 0;
  pushConstantRange.size = sizeof(PushConstantData);

  vk::PipelineLayoutCreateInfo pipelineLayoutInfo = {};
  pipelineLayoutInfo.setLayoutCount = 0;
  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");
  }
}

}  // namespace sophia