/*
 * Copyright 2022 Frederico de Oliveira Linhares
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "model.hpp"

#include <array>
#include <fstream>

#include "../command.hpp"
#include "../core.hpp"
#include "vertex.hpp"

namespace
{

// Data that is only needed for the command chain but not for the Model goes
// here.
struct ModelBuilder
{
  std::string model_path;
  VK::Model *model;

  ModelBuilder(VK::Model *m, std::string mp);
  ModelBuilder(VK::Model *m, const char* mp);
};

ModelBuilder::ModelBuilder(VK::Model *m, std::string mp):
  model{m},
  model_path{mp}
{
}

ModelBuilder::ModelBuilder(VK::Model *m, const char *mp):
  ModelBuilder{m, std::string(mp)}
{
}

uint32_t read_uint32_from_file(std::ifstream &input_file)
{
  uint32_t data{};
  input_file.read((char*)&data, sizeof(uint32_t));
  return data;
}

glm::vec2 read_vec2_from_file(std::ifstream &input_file)
{
  glm::vec2 data{};
  input_file.read((char*)&data.x, sizeof(glm::vec2::value_type));
  input_file.read((char*)&data.y, sizeof(glm::vec2::value_type));
  return data;
}

glm::vec3 read_vec3_from_file(std::ifstream &input_file)
{
  glm::vec3 data{};
  input_file.read((char*)&data.x, sizeof(glm::vec3::value_type));
  input_file.read((char*)&data.y, sizeof(glm::vec3::value_type));
  input_file.read((char*)&data.z, sizeof(glm::vec3::value_type));
  return data;
}

void
load_mesh(void *obj)
{
  auto self = static_cast<ModelBuilder*>(obj);

  std::ifstream input_file{self->model_path};
  if(!input_file.is_open()) throw CommandError{"Failed to open file."};

  self->model->queue_family =
    cg_core.vk_device_with_swapchain->get_queue_family_with_graphics();
  VK::Queue transfer_queue{self->model->queue_family->get_queue()};

  // Load vertexes.
  {
    auto vertex_count{read_uint32_from_file(input_file)};
    std::vector<VK::Vertex> vertexes{vertex_count};

    for(auto i{0}; i < vertex_count; i++)
    {
      vertexes[i].position = read_vec3_from_file(input_file);
      vertexes[i].normal = read_vec3_from_file(input_file);
      vertexes[i].texture_coord = read_vec2_from_file(input_file);
    }

    void *vertexes_data{vertexes.data()};
    size_t vertexes_size = sizeof(vertexes[0]) * vertexes.size();
    self->model->source_vertex_buffer = new VK::SourceBuffer{
      self->model->queue_family->device, vertexes_data, vertexes_size};
    self->model->vertex_buffer = new VK::DestinationBuffer{
      self->model->queue_family, self->model->source_vertex_buffer,
      VK_BUFFER_USAGE_VERTEX_BUFFER_BIT};
  }

  // Load indexes.
  {
    self->model->index_count = read_uint32_from_file(input_file);
    std::vector<uint32_t> indexes(self->model->index_count);

    for(auto i{0}; i < self->model->index_count; i++)
      indexes[i] = read_uint32_from_file(input_file);

    void *indexes_data{indexes.data()};
    size_t indexes_size{sizeof(indexes[0]) * indexes.size()};
    VK::SourceBuffer source_index_buffer{
      self->model->queue_family->device, indexes_data, indexes_size};
    self->model->index_buffer = new VK::DestinationBuffer{
      self->model->queue_family, &source_index_buffer,
      VK_BUFFER_USAGE_INDEX_BUFFER_BIT};
  }
}

void
unload_mesh(void *obj)
{
  auto self = static_cast<ModelBuilder*>(obj);

  delete self->model->index_buffer;
  delete self->model->vertex_buffer;
  delete self->model->source_vertex_buffer;
}

auto
load_uniform_buffers(void *obj)
{
  auto self = static_cast<ModelBuilder*>(obj);

  self->model->ub_model_instance.reserve(cg_core.vk_swapchain->images_count);
  for(int i{0}; i < cg_core.vk_swapchain->images_count; i++)
    self->model->ub_model_instance.emplace_back(
      cg_core.vk_device_with_swapchain, sizeof(VK::UBOModelInstance));
}

void
unload_uniform_buffers(void *obj)
{
  auto self = static_cast<ModelBuilder*>(obj);

  self->model->ub_model_instance.clear();
}

void
load_descriptor_set_pool(void *obj)
{
  auto self = static_cast<ModelBuilder*>(obj);

  std::array<VkDescriptorPoolSize, 2> descriptor_pool_sizes{};
  descriptor_pool_sizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  descriptor_pool_sizes[0].descriptorCount =
    self->model->ub_model_instance.size();
  descriptor_pool_sizes[1].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  descriptor_pool_sizes[1].descriptorCount =
    self->model->ub_model_instance.size();

  VkDescriptorPoolCreateInfo pool_info{};
  pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
  pool_info.pNext = nullptr;
  pool_info.flags = 0;
  pool_info.maxSets = self->model->ub_model_instance.size();
  pool_info.poolSizeCount = descriptor_pool_sizes.size();
  pool_info.pPoolSizes = descriptor_pool_sizes.data();

  if(vkCreateDescriptorPool(
       self->model->queue_family->device->device, &pool_info, nullptr,
       &self->model->descriptor_pool) != VK_SUCCESS)
    throw CommandError{"Failed to create a Vulkan descriptor pool."};
}

void
unload_descriptor_set_pool(void *obj)
{
  auto self = static_cast<ModelBuilder*>(obj);

  vkDestroyDescriptorPool(
    self->model->queue_family->device->device, self->model->descriptor_pool,
    nullptr);
}

void
load_descriptor_sets(void *obj)
{
  auto self = static_cast<ModelBuilder*>(obj);

  std::vector<VkDescriptorSetLayout> layouts(
    cg_core.vk_swapchain->images_count,
    cg_core.vk_graphics_pipeline->descriptor_set_layout_model_instance);

  VkDescriptorSetAllocateInfo alloc_info{};
  alloc_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
  alloc_info.descriptorPool = self->model->descriptor_pool;
  alloc_info.descriptorSetCount = layouts.size();
  alloc_info.pSetLayouts = layouts.data();

  self->model->descriptor_sets.resize(layouts.size());
  if(vkAllocateDescriptorSets(
       self->model->queue_family->device->device, &alloc_info,
       self->model->descriptor_sets.data()) != VK_SUCCESS)
    CommandError{"Failed to create Vulkan descriptor set."};
}

void
load_buffers_to_descriptor_sets(void *obj)
{
  auto self = static_cast<ModelBuilder*>(obj);

  for(auto i{0}; i < self->model->ub_model_instance.size(); i++)
  {
    VkDescriptorBufferInfo buffer_info{};
    buffer_info.buffer = self->model->ub_model_instance[i].buffer;
    buffer_info.offset = 0;
    buffer_info.range = sizeof(VK::UBOModelInstance);

    VkDescriptorImageInfo image_info{};
    image_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    image_info.imageView = self->model->texture->view;
    image_info.sampler = self->model->texture->sampler;

    std::array<VkWriteDescriptorSet, 2> write_descriptors{};
    write_descriptors[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write_descriptors[0].dstSet = self->model->descriptor_sets[i];
    write_descriptors[0].dstBinding = 0;
    write_descriptors[0].dstArrayElement = 0;
    write_descriptors[0].descriptorCount = 1;
    write_descriptors[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    write_descriptors[0].pBufferInfo = &buffer_info;
    write_descriptors[0].pImageInfo = nullptr;
    write_descriptors[0].pTexelBufferView = nullptr;

    write_descriptors[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write_descriptors[1].dstSet = self->model->descriptor_sets[i];
    write_descriptors[1].dstBinding = 1;
    write_descriptors[1].dstArrayElement = 0;
    write_descriptors[1].descriptorCount = 1;
    write_descriptors[1].descriptorType =
      VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    write_descriptors[1].pBufferInfo = nullptr;
    write_descriptors[1].pImageInfo = &image_info;
    write_descriptors[1].pTexelBufferView = nullptr;

    vkUpdateDescriptorSets(
      cg_core.vk_device_with_swapchain->device, write_descriptors.size(),
      write_descriptors.data(), 0, nullptr);
  }
}

static const CommandChain loader{
  {&load_mesh, &unload_mesh},
  {&load_uniform_buffers, &unload_uniform_buffers},
  {&load_descriptor_set_pool, &unload_descriptor_set_pool},
  {&load_descriptor_sets, nullptr},
  {&load_buffers_to_descriptor_sets, nullptr},
};

}

namespace VK
{

Model::Model(std::string model_path, std::shared_ptr<Texture> texture):
  texture{texture}
{
  ModelBuilder model_builder(this, model_path);
  loader.execute(&model_builder);
}

Model::Model(const char* model_path, std::shared_ptr<Texture> texture):
  Model{std::string(model_path), texture}
{
}

Model::~Model()
{
  ModelBuilder model_builder(this, "");
  loader.revert(&model_builder);
}

}