/*
 * Copyright 2022-2023 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 "skeletal_model.hpp"

#include "../core.hpp"
#include "uniform_data_object.hpp"

namespace
{

void
load_uniform_buffers(void *obj)
{
  auto self = static_cast<VK::SkeletalModel*>(obj);

  try
  {
    self->uniform_buffers.reserve(cg_core.vk_swapchain->images_count);
    for(auto i{0}; i < cg_core.vk_swapchain->images_count; i++)
      self->uniform_buffers.emplace_back(
	cg_core.vk_device_with_swapchain, sizeof(VK::UDOSkeletalModel));
  }
  catch(const std::exception& e)
  {
    throw CommandError{e.what()};
  }
}

void
unload_uniform_buffers(void *obj)
{
  auto self = static_cast<VK::SkeletalModel*>(obj);

  self->uniform_buffers.clear();
}

void
load_descriptor_set_pool(void *obj)
{
  auto self = static_cast<VK::SkeletalModel*>(obj);

  std::array<VkDescriptorPoolSize, 1> descriptor_pool_sizes{};
  descriptor_pool_sizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  descriptor_pool_sizes[0].descriptorCount =
    self->uniform_buffers.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->uniform_buffers.size();
  pool_info.poolSizeCount = descriptor_pool_sizes.size();
  pool_info.pPoolSizes = descriptor_pool_sizes.data();

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

void
unload_descriptor_set_pool(void *obj)
{
  auto self = static_cast<VK::SkeletalModel*>(obj);

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

void
load_descriptor_sets(void *obj)
{
  auto self = static_cast<VK::SkeletalModel*>(obj);

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

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

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

void
load_buffers_to_descriptor_sets(void *obj)
{
  auto self = static_cast<VK::SkeletalModel*>(obj);

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

    std::array<VkWriteDescriptorSet, 1> write_descriptors{};
    write_descriptors[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write_descriptors[0].dstSet = self->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;

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

static const CommandChain loader{
  {&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
{

SkeletalModel::SkeletalModel(
  std::shared_ptr<SkeletalMesh> skeletal_mesh,
  std::shared_ptr<Texture> texture, std::shared_ptr<glm::vec3> position,
  std::shared_ptr<glm::vec3> rotation):
  skeletal_mesh{skeletal_mesh},
  texture{texture},
  position{position},
  rotation{rotation},
  animation_index{0},
  animation_time{0.0f},
  bone_transforms(SKELETAL_MESH_MAX_NUM_OF_BONES)
{
  loader.execute(this);

  for(int i{0}; i < skeletal_mesh->bones.size(); i++)
    this->bone_transforms[i] = skeletal_mesh->bones[i].offset_matrix;
}

SkeletalModel::~SkeletalModel()
{
  loader.revert(this);
}

void
SkeletalModel::tick(float delta)
{
  VK::Animation *current_animation =
    &this->skeletal_mesh->animations[this->animation_index];

  { // update time
    this->animation_time += delta;
    if(this->animation_time > current_animation->final_time)
    {
      this->animation_time -= current_animation->final_time;
      for(VK::BoneTransform &bone_transform:
	    current_animation->bone_transforms)
      {
	bone_transform.positions.current_index = 0;
	bone_transform.rotations.current_index = 0;
	bone_transform.scales.current_index = 0;
      }
    }
  }

  for(int i{0}; i < current_animation->bone_transforms.size(); i++)
  {
    VK::BoneTransform *bone_transform = &current_animation->bone_transforms[i];

    auto position{bone_transform->positions.interpolate(
	this->animation_time,
	[](glm::vec3 frame)
	  {
	    return glm::translate(glm::mat4(1.0f), frame);
	  },
	[](glm::vec3 previous_frame, glm::vec3 next_frame, float scale_factor)
	  {
	    glm::vec3 final_position{glm::mix(
		previous_frame, next_frame, scale_factor)};
	    return glm::translate(glm::mat4(1.0f), final_position);
	  })};

    auto rotation{bone_transform->rotations.interpolate(
	this->animation_time,
	[](glm::quat frame)
	  {
	    return glm::toMat4(glm::normalize(frame));
	  },
	[](glm::quat previous_frame, glm::quat next_frame, float scale_factor)
	  {
	    return glm::toMat4(glm::slerp(
				 previous_frame, next_frame, scale_factor));
	  })};

    auto scale{bone_transform->scales.interpolate(
	this->animation_time,
	[](glm::vec3 frame)
	  {
	    return glm::scale(glm::mat4(1.0f), frame);
	  },
	[](glm::vec3 previous_frame, glm::vec3 next_frame, float scale_factor)
	  {
	    glm::vec3 scale{glm::mix(
		previous_frame, next_frame, scale_factor)};
	    return glm::scale(glm::mat4(1.0f), scale);
	  })};

    this->bone_transforms[i] = position * rotation * scale;
  }
}

}