#include <vulkan/vulkan.h>

#include "vulkan_engine/asset/object/model.h"
#include "vulkan_engine/vulkan/graphics.h"
#include "vulkan_engine/vulkan/uniform_transformations.h"

namespace veng {

std::uint32_t Graphics::FindMemoryType(
    std::uint32_t type_bits_filter, VkMemoryPropertyFlags required_properties) {
  VkPhysicalDeviceMemoryProperties memory_properties;
  vkGetPhysicalDeviceMemoryProperties(physical_device_, &memory_properties);
  gsl::span<VkMemoryType> memory_types(memory_properties.memoryTypes,
                                       memory_properties.memoryTypeCount);

  for (std::uint32_t i = 0; i < memory_types.size(); i++) {
    bool passes_filter = type_bits_filter & (1 << i);
    bool has_property_flags =
        memory_types[i].propertyFlags & required_properties;

    if (passes_filter && has_property_flags) return i;
  }

  throw std::runtime_error("Cannot find memory type!");
}

BufferHandle Graphics::CreateBuffer(VkDeviceSize size, VkBufferUsageFlags usage,
                                    VkMemoryPropertyFlags properties) {
  BufferHandle handle = {};

  VkBufferCreateInfo buffer_info = {};
  buffer_info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
  buffer_info.size = size;
  buffer_info.usage = usage;
  buffer_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

  VkResult result =
      vkCreateBuffer(logical_device_, &buffer_info, nullptr, &handle.buffer);
  if (result != VK_SUCCESS)
    throw std::runtime_error("Failed to create vertex buffer!");

  VkMemoryRequirements memory_requirements;
  vkGetBufferMemoryRequirements(logical_device_, handle.buffer,
                                &memory_requirements);

  std::uint32_t chosen_memory_type =
      FindMemoryType(memory_requirements.memoryTypeBits, properties);

  VkMemoryAllocateInfo allocation_info = {};
  allocation_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
  allocation_info.allocationSize = memory_requirements.size;
  allocation_info.memoryTypeIndex = chosen_memory_type;

  VkResult allocation_result = vkAllocateMemory(
      logical_device_, &allocation_info, nullptr, &handle.memory);
  if (allocation_result != VK_SUCCESS)
    throw std::runtime_error("Failed to allocate buffer memory!");

  vkBindBufferMemory(logical_device_, handle.buffer, handle.memory, 0);

  return handle;
}

BufferHandle Graphics::CreateVertexBuffer(gsl::span<Vertex> vertices) {
  VkDeviceSize size = sizeof(Vertex) * vertices.size();
  BufferHandle staging_handle = CreateBuffer(
      size,
      VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
          VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

  void* data;
  vkMapMemory(logical_device_, staging_handle.memory, 0, size, 0, &data);
  std::memcpy(data, vertices.data(), size);
  vkUnmapMemory(logical_device_, staging_handle.memory);

  BufferHandle gpu_handle = CreateBuffer(
      size,
      VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
      VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

  VkCommandBuffer transient_commands = BeginTransientCommandBuffer();

  VkBufferCopy copy_info = {};
  copy_info.srcOffset = 0;
  copy_info.dstOffset = 0;
  copy_info.size = size;
  vkCmdCopyBuffer(transient_commands, staging_handle.buffer, gpu_handle.buffer,
                  1, &copy_info);

  EndTransientCommandBuffer(transient_commands);

  DestroyBuffer(staging_handle);

  return gpu_handle;
}

BufferHandle Graphics::CreateIndexBuffer(gsl::span<std::uint32_t> indices) {
  VkDeviceSize size = sizeof(std::uint32_t) * indices.size();
  BufferHandle staging_handle =
      CreateBuffer(size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                   VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                       VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

  void* data;
  vkMapMemory(logical_device_, staging_handle.memory, 0, size, 0, &data);
  std::memcpy(data, indices.data(), size);
  vkUnmapMemory(logical_device_, staging_handle.memory);

  BufferHandle gpu_handle = CreateBuffer(
      size, VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
      VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

  VkCommandBuffer transient_commands = BeginTransientCommandBuffer();

  VkBufferCopy copy_info = {};
  copy_info.srcOffset = 0;
  copy_info.dstOffset = 0;
  copy_info.size = size;
  vkCmdCopyBuffer(transient_commands, staging_handle.buffer, gpu_handle.buffer,
                  1, &copy_info);

  EndTransientCommandBuffer(transient_commands);

  DestroyBuffer(staging_handle);

  return gpu_handle;
}

void Graphics::DestroyBuffer(BufferHandle handle) {
  vkDeviceWaitIdle(logical_device_);
  vkDestroyBuffer(logical_device_, handle.buffer, nullptr);
  vkFreeMemory(logical_device_, handle.memory, nullptr);
}

void Graphics::SetModelMatrix(glm::mat4 model) {
  // for (Frame& frame : frames_)
  vkCmdPushConstants(frames_[current_frame_].command_buffer, pipeline_layout_,
                     VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(glm::mat4), &model);
}

void Graphics::SetViewProjection(glm::mat4 view, glm::mat4 projection) {
  UniformTransformations transformations{view, projection};
  for (Frame& frame : frames_)
    std::memcpy(frame.uniform_buffer_location, &transformations,
                sizeof(UniformTransformations));
}

void Graphics::RenderBuffer(BufferHandle handle, std::uint32_t vertex_count) {
  VkDeviceSize offset = 0;
  vkCmdBindDescriptorSets(frames_[current_frame_].command_buffer,
                          VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_layout_, 0,
                          1, &frames_[current_frame_].uniform_set, 0, nullptr);
  vkCmdBindVertexBuffers(frames_[current_frame_].command_buffer, 0, 1,
                         &handle.buffer, &offset);
  vkCmdDraw(frames_[current_frame_].command_buffer, vertex_count, 1, 0, 0);
}

void Graphics::RenderIndexedBuffer(BufferHandle vertex_buffer,
                                   BufferHandle index_buffer,
                                   std::uint32_t count) {
  VkDeviceSize offset = 0;
  vkCmdBindDescriptorSets(frames_[current_frame_].command_buffer,
                          VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_layout_, 0,
                          1, &frames_[current_frame_].uniform_set, 0, nullptr);
  vkCmdBindVertexBuffers(frames_[current_frame_].command_buffer, 0, 1,
                         &vertex_buffer.buffer, &offset);
  vkCmdBindIndexBuffer(frames_[current_frame_].command_buffer,
                       index_buffer.buffer, 0, VK_INDEX_TYPE_UINT32);
  vkCmdDrawIndexed(frames_[current_frame_].command_buffer, count, 1, 0, 0, 0);
  SetModelMatrix(glm::mat4(1.f));
}
void Graphics::RenderModel(std::shared_ptr<Model> model) {
  if (!model->visible) return;
  SetTexture(model->material.texture_handle);
  SetModelMatrix(model->transform);
  RenderIndexedBuffer(model->vertex_buffer, model->index_buffer,
                      model->indices.size());
}

VkCommandBuffer Graphics::BeginTransientCommandBuffer() {
  VkCommandBufferAllocateInfo allocation_info = {};
  allocation_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
  allocation_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
  allocation_info.commandPool = command_pool_;
  allocation_info.commandBufferCount = 1;

  VkCommandBuffer buffer;
  vkAllocateCommandBuffers(logical_device_, &allocation_info, &buffer);

  VkCommandBufferBeginInfo begin_info = {};
  begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
  begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
  vkBeginCommandBuffer(buffer, &begin_info);

  return buffer;
}

void Graphics::EndTransientCommandBuffer(VkCommandBuffer command_buffer) {
  vkEndCommandBuffer(command_buffer);

  VkSubmitInfo submit_info = {};
  submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
  submit_info.commandBufferCount = 1;
  submit_info.pCommandBuffers = &command_buffer;

  vkQueueSubmit(graphics_queue_, 1, &submit_info, VK_NULL_HANDLE);
  vkQueueWaitIdle(graphics_queue_);
  vkFreeCommandBuffers(logical_device_, command_pool_, 1, &command_buffer);
}

void Graphics::CreateUniformBuffers() {
  VkDeviceSize buffer_size = sizeof(UniformTransformations);
  for (Frame& frame : frames_) {
    frame.uniform_buffer =
        CreateBuffer(buffer_size, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
                     VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                         VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

    vkMapMemory(logical_device_, frame.uniform_buffer.memory, 0, buffer_size, 0,
                &frame.uniform_buffer_location);
  }
}

void Graphics::CreateDescriptorSetLayouts() {
  VkDescriptorSetLayoutBinding uniform_layout_binding = {};
  uniform_layout_binding.binding = 0;
  uniform_layout_binding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  uniform_layout_binding.descriptorCount = 1;
  uniform_layout_binding.stageFlags = VK_SHADER_STAGE_ALL_GRAPHICS;

  VkDescriptorSetLayoutCreateInfo uniform_layout_info = {};
  uniform_layout_info.sType =
      VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
  uniform_layout_info.bindingCount = 1;
  uniform_layout_info.pBindings = &uniform_layout_binding;

  if (vkCreateDescriptorSetLayout(logical_device_, &uniform_layout_info,
                                  nullptr,
                                  &uniform_set_layout_) != VK_SUCCESS) {
    std::exit(EXIT_FAILURE);
  }

  VkDescriptorSetLayoutBinding texture_layout_binding = {};
  texture_layout_binding.binding = 0;
  texture_layout_binding.descriptorType =
      VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  texture_layout_binding.descriptorCount = 1;
  texture_layout_binding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

  VkDescriptorSetLayoutCreateInfo texture_layout_info = {};
  texture_layout_info.sType =
      VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
  texture_layout_info.bindingCount = 1;
  texture_layout_info.pBindings = &texture_layout_binding;

  if (vkCreateDescriptorSetLayout(logical_device_, &texture_layout_info,
                                  nullptr,
                                  &texture_set_layout_) != VK_SUCCESS) {
    std::exit(EXIT_FAILURE);
  }
}

void Graphics::CreateDescriptorPools() {
  VkDescriptorPoolSize uniform_pool_size = {};
  uniform_pool_size.type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  uniform_pool_size.descriptorCount = MAX_BUFFERED_FRAMES;

  VkDescriptorPoolCreateInfo uniform_create_info = {};
  uniform_create_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
  uniform_create_info.poolSizeCount = 1;
  uniform_create_info.pPoolSizes = &uniform_pool_size;
  uniform_create_info.maxSets = MAX_BUFFERED_FRAMES;

  if (vkCreateDescriptorPool(logical_device_, &uniform_create_info, nullptr,
                             &uniform_pool_) != VK_SUCCESS) {
    std::exit(EXIT_FAILURE);
  }

  VkPhysicalDeviceProperties properties = {};
  vkGetPhysicalDeviceProperties(physical_device_, &properties);

  VkDescriptorPoolSize texture_pool_size = {};
  texture_pool_size.type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  texture_pool_size.descriptorCount =
      properties.limits.maxSamplerAllocationCount;

  VkDescriptorPoolCreateInfo texture_create_info = {};
  texture_create_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
  texture_create_info.poolSizeCount = 1;
  texture_create_info.pPoolSizes = &texture_pool_size;
  texture_create_info.maxSets = properties.limits.maxSamplerAllocationCount;
  texture_create_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;

  if (vkCreateDescriptorPool(logical_device_, &texture_create_info, nullptr,
                             &texture_pool_) != VK_SUCCESS) {
    std::exit(EXIT_FAILURE);
  }
}

void Graphics::CreateDescriptorSets() {
  for (Frame& frame : frames_) {
    VkDescriptorSetAllocateInfo set_info = {};
    set_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    set_info.descriptorPool = uniform_pool_;
    set_info.descriptorSetCount = 1;
    set_info.pSetLayouts = &uniform_set_layout_;

    VkResult result = vkAllocateDescriptorSets(logical_device_, &set_info,
                                               &frame.uniform_set);
    if (result != VK_SUCCESS) std::exit(EXIT_FAILURE);

    VkDescriptorBufferInfo buffer_info = {};
    buffer_info.buffer = frame.uniform_buffer.buffer;
    buffer_info.offset = 0;
    buffer_info.range = sizeof(UniformTransformations);

    VkWriteDescriptorSet descriptor_write = {};
    descriptor_write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptor_write.dstSet = frame.uniform_set;
    descriptor_write.dstBinding = 0;
    descriptor_write.dstArrayElement = 0;
    descriptor_write.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    descriptor_write.descriptorCount = 1;
    descriptor_write.pBufferInfo = &buffer_info;

    vkUpdateDescriptorSets(logical_device_, 1, &descriptor_write, 0, nullptr);
  }
}

}  // namespace veng