#include <vulkan/vulkan.h>

#include "stb/stb_image.h"
#include "vulkan/graphics.h"

namespace veng {

void Graphics::CreateTextureSampler() {
  VkSamplerCreateInfo sampler_info = {};
  sampler_info.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
  sampler_info.magFilter = VK_FILTER_LINEAR;
  sampler_info.minFilter = VK_FILTER_LINEAR;
  sampler_info.addressModeU = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE;
  sampler_info.addressModeV = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE;
  sampler_info.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
  sampler_info.anisotropyEnable = VK_FALSE;
  sampler_info.maxAnisotropy = 1.f;
  sampler_info.borderColor = VK_BORDER_COLOR_FLOAT_OPAQUE_BLACK;
  sampler_info.unnormalizedCoordinates = VK_FALSE;
  sampler_info.compareEnable = VK_FALSE;
  sampler_info.compareOp = VK_COMPARE_OP_ALWAYS;
  sampler_info.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
  sampler_info.mipLodBias = 0.f;
  sampler_info.minLod = 0.f;
  sampler_info.maxLod = 0.f;

  if (vkCreateSampler(logical_device_, &sampler_info, nullptr,
                      &texture_sampler_) != VK_SUCCESS) {
    std::exit(EXIT_FAILURE);
  }
}

TextureHandle Graphics::CreateTexture(gsl::czstring path) {
  std::vector<std::uint8_t> data = ReadFile(path);
  return CreateTexture({data.data(), data.size()});
}

TextureHandle Graphics::CreateTexture(
    std::vector<std::uint8_t> image_file_data) {
  return CreateTexture({image_file_data.data(), image_file_data.size()});
}

TextureHandle Graphics::CreateTexture(gsl::span<std::uint8_t> image_file_data) {
  glm::ivec2 image_extents;
  std::int32_t channels;
  stbi_uc* pixel_data = stbi_load_from_memory(
      image_file_data.data(), image_file_data.size(), &image_extents.x,
      &image_extents.y, &channels, STBI_rgb_alpha);

  VkDeviceSize buffer_size = image_extents.x * image_extents.y * 4;
  BufferHandle staging =
      CreateBuffer(buffer_size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                   VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                       VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
  void* data_location;
  vkMapMemory(logical_device_, staging.memory, 0, buffer_size, 0,
              &data_location);
  std::memcpy(data_location, pixel_data, buffer_size);
  vkUnmapMemory(logical_device_, staging.memory);

  stbi_image_free(pixel_data);

  TextureHandle handle =
      CreateImage(image_extents, VK_FORMAT_R8G8B8A8_SRGB,
                  VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
                  VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

  TransitionImageLayout(handle.image, VK_IMAGE_LAYOUT_UNDEFINED,
                        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
  CopyBufferToImage(staging.buffer, handle.image, image_extents);
  TransitionImageLayout(handle.image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

  handle.image_view = CreateImageView(handle.image, VK_FORMAT_R8G8B8A8_SRGB,
                                      VK_IMAGE_ASPECT_COLOR_BIT);

  VkDescriptorSetAllocateInfo set_info = {};
  set_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
  set_info.descriptorPool = texture_pool_;
  set_info.descriptorSetCount = 1;
  set_info.pSetLayouts = &texture_set_layout_;

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

  VkDescriptorImageInfo image_info = {};
  image_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
  image_info.imageView = handle.image_view;
  image_info.sampler = texture_sampler_;

  VkWriteDescriptorSet descriptor_write = {};
  descriptor_write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
  descriptor_write.dstSet = handle.set;
  descriptor_write.dstBinding = 0;
  descriptor_write.dstArrayElement = 0;
  descriptor_write.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  descriptor_write.descriptorCount = 1;
  descriptor_write.pImageInfo = &image_info;

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

  DestroyBuffer(staging);
  return handle;
}

void Graphics::DestroyTexture(TextureHandle handle) {
  vkDeviceWaitIdle(logical_device_);

  if (handle.set != VK_NULL_HANDLE)
    vkFreeDescriptorSets(logical_device_, texture_pool_, 1, &handle.set);

  if (handle.image_view != VK_NULL_HANDLE)
    vkDestroyImageView(logical_device_, handle.image_view, nullptr);

  if (handle.image != VK_NULL_HANDLE)
    vkDestroyImage(logical_device_, handle.image, nullptr);

  if (handle.memory != VK_NULL_HANDLE)
    vkFreeMemory(logical_device_, handle.memory, nullptr);
}

void Graphics::SetTexture(TextureHandle handle) {
  vkCmdBindDescriptorSets(frames_[current_frame_].command_buffer,
                          VK_PIPELINE_BIND_POINT_GRAPHICS,
                          pipeline_layout_, 1, 1, &handle.set, 0, nullptr);
}

void Graphics::TransitionImageLayout(VkImage image, VkImageLayout old_layout,
                                     VkImageLayout new_layout) {
  VkCommandBuffer local_command_buffer = BeginTransientCommandBuffer();

  VkImageMemoryBarrier barrier = {};
  barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
  barrier.oldLayout = old_layout;
  barrier.newLayout = new_layout;
  barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.image = image;
  barrier.subresourceRange.baseArrayLayer = 0;
  barrier.subresourceRange.baseMipLevel = 0;
  barrier.subresourceRange.levelCount = 1;
  barrier.subresourceRange.layerCount = 1;

  VkPipelineStageFlags source_stage = {};
  VkPipelineStageFlags destination_stage = {};

  if (old_layout == VK_IMAGE_LAYOUT_UNDEFINED &&
      new_layout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {
    barrier.srcAccessMask = 0;
    barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    source_stage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
    destination_stage = VK_PIPELINE_STAGE_TRANSFER_BIT;
  } else if (old_layout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL &&
             new_layout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {
    barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
    source_stage = VK_PIPELINE_STAGE_TRANSFER_BIT;
    destination_stage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
  } else if (old_layout == VK_IMAGE_LAYOUT_UNDEFINED &&
             new_layout == VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL) {
    barrier.srcAccessMask = 0;
    barrier.dstAccessMask = VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT |
                            VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
    source_stage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
    destination_stage = VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
  }

  if (new_layout == VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL)
    barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;
  else
    barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;

  vkCmdPipelineBarrier(local_command_buffer, source_stage, destination_stage, 0,
                       0, nullptr, 0, nullptr, 1, &barrier);

  EndTransientCommandBuffer(local_command_buffer);
}

void Graphics::CopyBufferToImage(VkBuffer buffer, VkImage image,
                                 glm::ivec2 image_size) {
  VkCommandBuffer local_command_buffer = BeginTransientCommandBuffer();

  VkBufferImageCopy region = {};
  region.bufferOffset = 0;
  region.bufferRowLength = 0;
  region.bufferImageHeight = 0;
  region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
  region.imageSubresource.mipLevel = 0;
  region.imageSubresource.baseArrayLayer = 0;
  region.imageSubresource.layerCount = 1;
  region.imageOffset = {0, 0, 0};
  region.imageExtent = {static_cast<std::uint32_t>(image_size.x),
                        static_cast<std::uint32_t>(image_size.y), 1};

  vkCmdCopyBufferToImage(local_command_buffer, buffer, image,
                         VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);

  EndTransientCommandBuffer(local_command_buffer);
}

TextureHandle Graphics::CreateImage(glm::ivec2 size, VkFormat image_format,
                                    VkBufferUsageFlags usage,
                                    VkMemoryPropertyFlags properties) {
  TextureHandle handle = {};

  VkImageCreateInfo image_info = {};
  image_info.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
  image_info.usage = usage;
  image_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
  image_info.imageType = VK_IMAGE_TYPE_2D;
  image_info.extent.width = size.x;
  image_info.extent.height = size.y;
  image_info.extent.depth = 1;
  image_info.mipLevels = 1;
  image_info.arrayLayers = 1;
  image_info.format = image_format;
  image_info.tiling = VK_IMAGE_TILING_OPTIMAL;
  image_info.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
  image_info.samples = VK_SAMPLE_COUNT_1_BIT;
  image_info.flags = 0;

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

  VkMemoryRequirements memory_requirements;
  vkGetImageMemoryRequirements(logical_device_, handle.image,
                               &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 image memory!");

  vkBindImageMemory(logical_device_, handle.image, handle.memory, 0);

  return handle;
}

void Graphics::CreateDepthResources() {
  VkFormat kDepthFormat = VK_FORMAT_D32_SFLOAT;
  depth_texture_ = CreateImage({extent_.width, extent_.height}, kDepthFormat,
                               VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT,
                               VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
  depth_texture_.image_view = CreateImageView(
      depth_texture_.image, kDepthFormat, VK_IMAGE_ASPECT_DEPTH_BIT);
}

}  // namespace veng