#include "resources/mesh.hpp"

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

namespace sophia {

std::vector<vk::VertexInputBindingDescription> Mesh::getBindingDescriptions() {
  return {{0, sizeof(MeshVertex), vk::VertexInputRate::eVertex}};
}

std::vector<vk::VertexInputAttributeDescription>
Mesh::getAttributeDescriptions() {
  return {
      {0, 0, vk::Format::eR32G32B32Sfloat, offsetof(MeshVertex, position)},
      {1, 0, vk::Format::eR32G32B32Sfloat, offsetof(MeshVertex, normal)},
      {2, 0, vk::Format::eR32G32Sfloat, offsetof(MeshVertex, uv)},
      {3, 0, vk::Format::eR32G32B32A32Sfloat, offsetof(MeshVertex, color)},
  };
}

Mesh::Mesh(Device& device, const MeshData& data) : device{device} {
  createVertexBuffers(data.vertices);
  createIndexBuffers(data.indices);
}

Mesh::~Mesh() = default;

void Mesh::bind(vk::CommandBuffer commandBuffer) const {
  vk::Buffer buffers[] = {this->vertexBuffer->getBuffer()};
  vk::DeviceSize offset[] = {0};
  commandBuffer.bindVertexBuffers(0, 1, buffers, offset);

  if (this->hasIndexBuffer) {
    commandBuffer.bindIndexBuffer(this->indexBuffer->getBuffer(), 0,
                                  vk::IndexType::eUint32);
  }
}

void Mesh::draw(vk::CommandBuffer commandBuffer, u32 instanceCount,
                u32 firstInstance) const {
  if (this->hasIndexBuffer) {
    commandBuffer.drawIndexed(this->indexCount, instanceCount, 0, 0,
                              firstInstance);
  } else {
    commandBuffer.draw(this->vertexCount, instanceCount, 0, firstInstance);
  }
}

void Mesh::createVertexBuffers(const std::vector<MeshVertex>& vertices) {
  this->vertexCount = static_cast<u32>(vertices.size());
  SOPHIA_ASSERT(this->vertexCount >= 3, "vertex count must be at least 3");

  vk::DeviceSize bufferSize = sizeof(vertices[0]) * this->vertexCount;

  u32 vertexSize = sizeof(vertices[0]);

  Buffer stagingBuffer{
      this->device,
      vertexSize,
      vertexCount,
      vk::BufferUsageFlagBits::eTransferSrc,
      vk::MemoryPropertyFlagBits::eHostVisible |
          vk::MemoryPropertyFlagBits::eHostCoherent,
  };

  stagingBuffer.map();
  stagingBuffer.writeToBuffer((void*)vertices.data());

  vertexBuffer =
      std::make_unique<Buffer>(this->device, vertexSize, vertexCount,
                               vk::BufferUsageFlagBits::eTransferDst |
                                   vk::BufferUsageFlagBits::eVertexBuffer,
                               vk::MemoryPropertyFlagBits::eDeviceLocal);

  this->device.copyBuffer(stagingBuffer.getBuffer(),
                          this->vertexBuffer->getBuffer(), bufferSize);
}

void Mesh::createIndexBuffers(const std::vector<u32>& indicies) {
  this->indexCount = static_cast<u32>(indicies.size());
  this->hasIndexBuffer = this->indexCount > 0;

  if (!this->hasIndexBuffer) {
    return;
  }

  vk::DeviceSize bufferSize = sizeof(indicies[0]) * this->indexCount;

  u32 indexSize = sizeof(indicies[0]);

  Buffer stagingBuffer{
      this->device,
      indexSize,
      indexCount,
      vk::BufferUsageFlagBits::eTransferSrc,
      vk::MemoryPropertyFlagBits::eHostVisible |
          vk::MemoryPropertyFlagBits::eHostCoherent,
  };

  stagingBuffer.map();
  stagingBuffer.writeToBuffer((void*)indicies.data());

  indexBuffer =
      std::make_unique<Buffer>(this->device, indexSize, indexCount,
                               vk::BufferUsageFlagBits::eTransferDst |
                                   vk::BufferUsageFlagBits::eIndexBuffer,
                               vk::MemoryPropertyFlagBits::eDeviceLocal);

  this->device.copyBuffer(stagingBuffer.getBuffer(),
                          this->indexBuffer->getBuffer(), bufferSize);
}

}  // namespace sophia