/*
 * Copyright 2022-2024 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 "texture.hpp"

#include "../command.hpp"
#include "../core.hpp"
#include "image.hpp"
#include "qoi.hpp"
#include "source_buffer.hpp"

namespace
{

inline void
create_vulkan_image(
  VkImage *image, VkDeviceMemory *device_memory, int width, int height,
  uint32_t mip_levels)
{
  VkExtent3D vk_extent3d{};
  vk_extent3d.width = width;
  vk_extent3d.height = height;
  vk_extent3d.depth = 1;

  VK::Image::create(
    cg_core.vk_device_with_swapchain,
    image,
    device_memory,
    VK_FORMAT_R8G8B8A8_UNORM,
    vk_extent3d,
    mip_levels,
    VK_IMAGE_TILING_OPTIMAL,
    VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT);
}

struct ImageBuilder
{
  VK::Texture *texture;
};

struct ImageTextureBuilder: public ImageBuilder
{
  std::string texture_path;

  ImageTextureBuilder(VK::Texture *t, std::string tp);
  ImageTextureBuilder(VK::Texture *t, const char* tp);
};

ImageTextureBuilder::ImageTextureBuilder(VK::Texture *t, std::string tp):
  texture_path{tp}
{
  this->texture = t;
}

ImageTextureBuilder::ImageTextureBuilder(VK::Texture *t, const char* tp):
  ImageTextureBuilder{t, std::string(tp)}
{
}

void
load_image(void *obj)
{
  auto self = static_cast<ImageTextureBuilder*>(obj);

  const int num_channels = 4; // all images are converted to RGBA
  VK::QOI::Image qoi_image(self->texture_path.c_str(), num_channels);
  uint8_t *pixels;

  { // Load file image from file.
    self->texture->width = qoi_image.header.width;
    self->texture->height = qoi_image.header.height;
    self->texture->mip_levels = 1;

    pixels = qoi_image.pixels;
  }

  // Load file image into a vulkan buffer.
  size_t image_size{static_cast<size_t>(
      qoi_image.header.width * qoi_image.header.height * num_channels)};
  VK::SourceBuffer source_image_buffer{
    cg_core.vk_device_with_swapchain, pixels, image_size};

  { // Create vulkan image.
    try
    {
      create_vulkan_image(
        &self->texture->image,
        &self->texture->device_memory,
        self->texture->width,
        self->texture->height,
        self->texture->mip_levels);
    }
    catch(VK::Image::Error error)
    {
      throw CommandError{error.what()};
    }
  }

  // Copy image from vulkan buffer into vulkan image.
  {
    auto queue_family = self->texture->queue_family;
    auto queue{queue_family->get_queue()};
    VK::CommandPool command_pool{queue_family, 1};
    VkCommandBuffer vk_command_buffer{command_pool.command_buffers[0]};

    queue.submit_one_time_command(vk_command_buffer, [&](){
      VK::Image::move_image_state(
	vk_command_buffer, self->texture->image, VK_FORMAT_R8G8B8A8_UNORM,
	0, VK_ACCESS_TRANSFER_WRITE_BIT,
	VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
	VK_PIPELINE_STAGE_HOST_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);

      VkBufferImageCopy image_copy{};
      image_copy.bufferOffset = 0;
      image_copy.bufferRowLength = 0;
      image_copy.bufferImageHeight = 0;
      image_copy.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
      image_copy.imageSubresource.mipLevel = 0;
      image_copy.imageSubresource.baseArrayLayer = 0;
      image_copy.imageSubresource.layerCount = 1;
      image_copy.imageOffset = {0, 0, 0};
      image_copy.imageExtent = {
        self->texture->width, self->texture->height, 1};

      vkCmdCopyBufferToImage(
	vk_command_buffer, source_image_buffer.buffer, self->texture->image,
	VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &image_copy);

      VK::Image::move_image_state(
	vk_command_buffer, self->texture->image, VK_FORMAT_R8G8B8A8_UNORM,
	VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT,
	VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
	VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
	VK_PIPELINE_STAGE_TRANSFER_BIT,
	VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT);
    });
  }
}

void
unload_image(void *obj)
{
  auto self = static_cast<ImageBuilder*>(obj);

  vkDestroyImage(
    cg_core.vk_device_with_swapchain->device, self->texture->image, nullptr);
  vkFreeMemory(
    cg_core.vk_device_with_swapchain->device, self->texture->device_memory,
    nullptr);
}

void
load_sampler(void *obj)
{
  auto self = static_cast<ImageBuilder*>(obj);

  VkSamplerCreateInfo sampler_info{};
  sampler_info.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
  sampler_info.pNext = nullptr;
  sampler_info.flags = 0;
  sampler_info.magFilter = VK_FILTER_LINEAR;
  sampler_info.minFilter = VK_FILTER_LINEAR;
  sampler_info.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
  sampler_info.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;
  sampler_info.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;
  sampler_info.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
  sampler_info.mipLodBias = 0.0f;
  sampler_info.anisotropyEnable = VK_TRUE;
  sampler_info.maxAnisotropy = 16;
  sampler_info.compareEnable = VK_FALSE;
  sampler_info.compareOp = VK_COMPARE_OP_NEVER;
  sampler_info.minLod = 0.0f;
  sampler_info.maxLod = 0.0f;
  sampler_info.borderColor = VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE;
  sampler_info.unnormalizedCoordinates = VK_FALSE;

  if(vkCreateSampler(
       cg_core.vk_device_with_swapchain->device, &sampler_info, nullptr,
       &self->texture->sampler) != VK_SUCCESS)
    throw CommandError{"Failed to create texture sampler."};
}

void
unload_sampler(void *obj)
{
  auto self = static_cast<ImageBuilder*>(obj);

  vkDestroySampler(
    cg_core.vk_device_with_swapchain->device, self->texture->sampler, nullptr);
}

void
load_view(void *obj)
{
  auto self = static_cast<ImageBuilder*>(obj);

  try
  {
    VK::Image::create_view(
      cg_core.vk_device_with_swapchain, &self->texture->view,
      self->texture->image,
      VK_FORMAT_R8G8B8A8_UNORM, VK_IMAGE_ASPECT_COLOR_BIT);
  }
  catch(VK::Image::Error error)
  {
    throw CommandError{error.what()};
  }
}

void
unload_view(void *obj)
{
  auto self = static_cast<ImageBuilder*>(obj);

  vkDestroyImageView(
    cg_core.vk_device_with_swapchain->device, self->texture->view, nullptr);
}

const CommandChain image_loader{
  {&load_image, &unload_image},
  {&load_sampler, &unload_sampler},
  {&load_view, &unload_view}
};

struct CharacterToDraw
{
  int pos_x;
  std::shared_ptr<VK::Character> character;

  CharacterToDraw(int x, std::shared_ptr<VK::Character> character):
    pos_x{x},
    character{character}
    {};
};

struct TextTextureBuilder: public ImageBuilder
{
  VK::Font *font;
  const char* str;
  uint32_t max_bearing_y;
  std::vector<CharacterToDraw> chars_to_draw;

  TextTextureBuilder(VK::Texture *texture, VK::Font *font, const char* str):
    font{font},
    str{str}
    {
      this->texture = texture;
    }
};

void
load_text_proportions(void *obj)
{
  auto self = static_cast<TextTextureBuilder*>(obj);

  uint32_t texture_width{0}, texture_descender{0};
  auto unicode_text{VK::Character::str_to_unicode(self->str)};

  auto first_image{self->font->character(unicode_text[0])};
  if(first_image->bearing_x < 0) texture_width = - first_image->bearing_x;

  self->max_bearing_y = 0;
  self->chars_to_draw.reserve(unicode_text.size());

  // FIXME: I need to test several different fonts to find all bugs in this
  // code.
  std::shared_ptr<VK::Character> char_image{};
  { // Calculate image size
    int max_height;
    for(auto char_code : unicode_text)
    {
      char_image = self->font->character(char_code);
      uint32_t descender{char_image->height - char_image->bearing_y};
      uint32_t pos_x{texture_width + char_image->bearing_x};

      if(char_image->image != VK_NULL_HANDLE)
        self->chars_to_draw.emplace_back(pos_x, char_image);

      if(char_image->bearing_y > self->max_bearing_y)
        self->max_bearing_y = char_image->bearing_y;
      if(descender > texture_descender) texture_descender = descender;

      texture_width += char_image->advance;
    }
  }

  { // Restore image width if last character have a negative bearing.
    int bearing_x_pluss_width = char_image->bearing_x + char_image->width;
    if(bearing_x_pluss_width > char_image->advance)
      texture_width += bearing_x_pluss_width - char_image->advance;
  }

  self->texture->width = texture_width;
  self->texture->height = self->max_bearing_y + texture_descender;
  self->texture->mip_levels = 1;
}

void
load_text_image(void *obj)
{
  auto self = static_cast<TextTextureBuilder*>(obj);

  const int NumChannels = 4;

  size_t image_size{static_cast<size_t>(
      self->texture->width * self->texture->height * NumChannels)};
  std::vector<unsigned char> pixels(image_size);
  for(auto x{0}; x < self->texture->width; x++)
  {
    for(auto y{0}; y < self->texture->height; y++)
    {
      auto image_coord = y * self->font->face->glyph->bitmap.width +
        x * NumChannels;
      pixels[image_coord] = 0; // Red
      pixels[image_coord + 1] = 0; // Green
      pixels[image_coord + 2] = 0; // Blue
      pixels[image_coord + 3] = 0; // Alpha
    }
  }
  VK::SourceBuffer source_image_buffer{
    cg_core.vk_device_with_swapchain, pixels.data(), image_size};

  { // Create vulkan image.
    try
    {
      create_vulkan_image(
        &self->texture->image,
        &self->texture->device_memory,
        self->texture->width,
        self->texture->height,
        self->texture->mip_levels);
    }
    catch(VK::Image::Error error)
    {
      throw CommandError{error.what()};
    }
  }

  { // Render text
    auto queue_family{
      cg_core.vk_device_with_swapchain->get_queue_family_with_presentation()};
    auto queue{queue_family->get_queue()};
    VK::CommandPool command_pool{queue_family, 1};
    VkCommandBuffer vk_command_buffer{command_pool.command_buffers[0]};

    queue.submit_one_time_command(vk_command_buffer, [&](){
      VK::Image::move_image_state(
	vk_command_buffer, self->texture->image, VK_FORMAT_R8G8B8A8_UNORM,
	0, VK_ACCESS_TRANSFER_WRITE_BIT,
	VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
	VK_PIPELINE_STAGE_HOST_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);

      VkBufferImageCopy image_copy{};
      image_copy.bufferOffset = 0;
      image_copy.bufferRowLength = 0;
      image_copy.bufferImageHeight = 0;
      image_copy.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
      image_copy.imageSubresource.mipLevel = 0;
      image_copy.imageSubresource.baseArrayLayer = 0;
      image_copy.imageSubresource.layerCount = 1;
      image_copy.imageOffset = {0, 0, 0};
      image_copy.imageExtent = {self->texture->width, self->texture->height, 1};

      vkCmdCopyBufferToImage(
	vk_command_buffer, source_image_buffer.buffer, self->texture->image,
	VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &image_copy);

      for(auto &to_draw: self->chars_to_draw)
      {
        VkImageSubresourceLayers
          source_subresources{}, destination_subresources{};
        source_subresources.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        source_subresources.mipLevel = 0;
        source_subresources.baseArrayLayer = 0;
        source_subresources.layerCount = 1;

        destination_subresources.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        destination_subresources.mipLevel = 0;
        destination_subresources.baseArrayLayer = 0;
        destination_subresources.layerCount = 1;

        VkImageCopy image_copy{};
        image_copy.srcSubresource = source_subresources;
        image_copy.srcOffset = {0, 0, 0};
        image_copy.dstSubresource = destination_subresources;
        image_copy.dstOffset = {
          to_draw.pos_x,
          (int)(self->max_bearing_y - to_draw.character->bearing_y),
          0};
        image_copy.extent = {
          to_draw.character->width, to_draw.character->height, 1};

        vkCmdCopyImage(
          vk_command_buffer,
          to_draw.character->image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
          self->texture->image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
          1, &image_copy);
      }

      VK::Image::move_image_state(
	vk_command_buffer, self->texture->image, VK_FORMAT_R8G8B8A8_UNORM,
	VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT,
	VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
	VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
	VK_PIPELINE_STAGE_TRANSFER_BIT,
	VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT);
    });
  }
}

const CommandChain text_loader{
  {&load_text_proportions, nullptr},
  {&load_text_image, &unload_image},
  {&load_sampler, &unload_sampler},
  {&load_view, &unload_view}
};

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

  std::array<VkDescriptorPoolSize, 1> descriptor_pool_sizes{};
  descriptor_pool_sizes[0].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  descriptor_pool_sizes[0].descriptorCount =
    cg_core.vk_swapchain->images_count;

  VkDescriptorPoolCreateInfo pool_info{};
  pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
  pool_info.pNext = nullptr;
  pool_info.flags = 0;
  pool_info.maxSets = cg_core.vk_swapchain->images_count;
  pool_info.poolSizeCount = descriptor_pool_sizes.size();
  pool_info.pPoolSizes = descriptor_pool_sizes.data();

  if(vkCreateDescriptorPool(
       self->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::Texture*>(obj);

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

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

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

  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->queue_family->device->device, &alloc_info,
       self->descriptor_sets.data()) != VK_SUCCESS)
    CommandError{"Failed to create Vulkan descriptor set."};
}

void
load_data_to_descriptor_sets(void *obj)
{
  auto self = static_cast<VK::Texture*>(obj);

  for(auto i{0}; i < cg_core.vk_swapchain->images_count; i++)
  {
    VkDescriptorImageInfo image_info{};
    image_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    image_info.imageView = self->view;
    image_info.sampler = self->sampler;

    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_COMBINED_IMAGE_SAMPLER;
    write_descriptors[0].pBufferInfo = nullptr;
    write_descriptors[0].pImageInfo = &image_info;
    write_descriptors[0].pTexelBufferView = nullptr;

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

const CommandChain descriptor_loader{
  {&load_descriptor_set_pool, &unload_descriptor_set_pool},
  {&load_descriptor_sets, nullptr},
  {&load_data_to_descriptor_sets, nullptr}
};

}

namespace VK
{

Texture::Texture(Font *font, const char* str)
{
  this->queue_family =
    cg_core.vk_device_with_swapchain->get_queue_family_with_presentation();

  TextTextureBuilder text_builder(this, font, str);
  text_loader.execute(&text_builder);
  descriptor_loader.execute(this);
}

Texture::Texture(std::string texture_path)
{
  this->queue_family =
    cg_core.vk_device_with_swapchain->get_queue_family_with_presentation();

  ImageTextureBuilder texture_builder(this, texture_path);
  image_loader.execute(&texture_builder);
  descriptor_loader.execute(this);
}

Texture::Texture(const char* texture_path):
  Texture{std::string(texture_path)}
{
}

Texture::~Texture()
{
  ImageTextureBuilder texture_builder(this, "");
  image_loader.revert(&texture_builder);
  descriptor_loader.revert(this);
}

}