#define STB_IMAGE_IMPLEMENTATION
#include <stdafx.h>
#include "Texture.h"
#include "Buffer/Buffer.h"
#include "Context.h"
#include "Device.h"
#include "RNG.h"



namespace fpr
{
vk::UniqueImageView CreateImageView(
    fpr::Device*         device,
    vk::Image            image,
    vk::Format           format,
    vk::ImageAspectFlags aspect_mask,
    vk::ImageViewType    view_type,
    uint32_t             mip_levels)
{
  vk::ImageViewCreateInfo image_view_create_info{};
  image_view_create_info.setImage(image);
  image_view_create_info.setViewType(view_type);
  image_view_create_info.format = format;

  image_view_create_info.components.setR(vk::ComponentSwizzle::eIdentity);
  image_view_create_info.components.setG(vk::ComponentSwizzle::eIdentity);
  image_view_create_info.components.setB(vk::ComponentSwizzle::eIdentity);
  image_view_create_info.components.setA(vk::ComponentSwizzle::eIdentity);

  image_view_create_info.subresourceRange.aspectMask     = aspect_mask;
  image_view_create_info.subresourceRange.baseMipLevel   = 0;
  image_view_create_info.subresourceRange.levelCount     = mip_levels;
  image_view_create_info.subresourceRange.baseArrayLayer = 0;
  image_view_create_info.subresourceRange.layerCount     = 1;
  auto [result_signal, image_view] = device->GetLogicalDeviceHandle().createImageViewUnique(image_view_create_info);
  assert(("Failed to create image view!", result_signal == vk::Result::eSuccess));
  return std::move(image_view);
}


std::pair<vk::DeviceSize, stbi_uc*> Texture::GetImageDetails(std::string_view file_name, int desired_channels)
{
  int      width;
  int      height;
  int      image_size;
  stbi_uc* raw_data = stbi_load(file_name.data(), &width, &height, &image_size, desired_channels);
  assert(raw_data != nullptr);

  m_width         = (uint32_t)width;
  m_height        = (uint32_t)height;
  m_channels_used = desired_channels;
  m_mip_levels    = static_cast<uint32_t>(std::floor(std::log2(std::max(m_width, m_height)))) + 1;
  return { vk::DeviceSize((vk::DeviceSize)width * (vk::DeviceSize)height * (vk::DeviceSize)desired_channels),
           raw_data };
}

const vk::DescriptorSet& Texture::GetDescriptorSet() const
{
  return *m_descriptor_set;
}

void Texture::GenerateMipMaps()
{
  auto& context    = fpr::Context::Get();
  auto& cmd_buffer = context.BeginSingleTimeCommands();

  auto half_else_one = [=](uint32_t to_half) { return to_half > 1 ? to_half / 2 : 1; };

  vk::ImageMemoryBarrier mip_img_barrier;
  mip_img_barrier.image                           = *m_image;
  mip_img_barrier.subresourceRange.aspectMask     = vk::ImageAspectFlagBits::eColor;
  mip_img_barrier.subresourceRange.layerCount     = 1;
  mip_img_barrier.subresourceRange.levelCount     = 1;
  mip_img_barrier.subresourceRange.baseArrayLayer = 0;
  mip_img_barrier.srcQueueFamilyIndex             = 0;
  mip_img_barrier.dstQueueFamilyIndex             = 0;

  uint32_t mip_width  = m_width;
  uint32_t mip_height = m_height;

  for(uint32_t i = 1; i < m_mip_levels; ++i)
  {
    mip_img_barrier.subresourceRange.baseMipLevel = i - 1;
    mip_img_barrier.oldLayout                     = vk::ImageLayout::eTransferDstOptimal;
    mip_img_barrier.newLayout                     = vk::ImageLayout::eTransferSrcOptimal;
    mip_img_barrier.srcAccessMask                 = vk::AccessFlagBits::eTransferWrite;
    mip_img_barrier.dstAccessMask                 = vk::AccessFlagBits::eTransferRead;

    cmd_buffer.pipelineBarrier(
        vk::PipelineStageFlagBits::eTransfer,
        vk::PipelineStageFlagBits::eTransfer,
        vk::DependencyFlagBits::eByRegion,
        nullptr,
        nullptr,
        mip_img_barrier);

    vk::ImageBlit blit;

    blit.srcOffsets[0]             = vk::Offset3D(0, 0, 0);
    blit.srcOffsets[1]             = vk::Offset3D(mip_width, mip_height, 1);
    blit.srcSubresource.aspectMask = vk::ImageAspectFlagBits::eColor;
    blit.srcSubresource.mipLevel   = i - 1;
    blit.srcSubresource.layerCount = 1;

    blit.dstOffsets[0]             = vk::Offset3D(0, 0, 0);
    blit.dstOffsets[1]             = vk::Offset3D(half_else_one(mip_width), half_else_one(mip_height), 1);
    blit.dstSubresource.aspectMask = vk::ImageAspectFlagBits::eColor;
    blit.dstSubresource.mipLevel   = i;
    blit.dstSubresource.layerCount = 1;

    cmd_buffer.blitImage(
        *m_image,
        vk::ImageLayout::eTransferSrcOptimal,
        *m_image,
        vk::ImageLayout::eTransferDstOptimal,
        blit,
        vk::Filter::eLinear);

    mip_img_barrier.oldLayout     = vk::ImageLayout::eTransferSrcOptimal;
    mip_img_barrier.newLayout     = vk::ImageLayout::eShaderReadOnlyOptimal;
    mip_img_barrier.srcAccessMask = vk::AccessFlagBits::eTransferRead;
    mip_img_barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead;

    cmd_buffer.pipelineBarrier(
        vk::PipelineStageFlagBits::eTransfer,
        vk::PipelineStageFlagBits::eFragmentShader,
        vk::DependencyFlagBits::eByRegion,
        nullptr,
        nullptr,
        mip_img_barrier);

    mip_width  = half_else_one(mip_width);
    mip_height = half_else_one(mip_height);
  }

  mip_img_barrier.subresourceRange.baseMipLevel = m_mip_levels - 1;
  mip_img_barrier.oldLayout                     = vk::ImageLayout::eTransferDstOptimal;
  mip_img_barrier.newLayout                     = vk::ImageLayout::eShaderReadOnlyOptimal;
  mip_img_barrier.srcAccessMask                 = vk::AccessFlagBits::eTransferWrite;
  mip_img_barrier.dstAccessMask                 = vk::AccessFlagBits::eShaderRead;

  cmd_buffer.pipelineBarrier(
      vk::PipelineStageFlagBits::eTransfer,
      vk::PipelineStageFlagBits::eFragmentShader,
      vk::DependencyFlagBits::eByRegion,
      nullptr,
      nullptr,
      mip_img_barrier);

  context.EndSingleTimeCommands();
}

Texture::Texture(std::string_view file_name, int desired_channels, vk::Format format): m_format(format)
{
  auto& context = fpr::Context::Get();
  auto  device  = context.GetDevice();

  auto [size, raw_data]         = GetImageDetails(file_name, desired_channels);
  vk::MemoryPropertyFlags flags = vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent;
  fpr::Buffer             staging_buffer(size, vk::BufferUsageFlagBits::eTransferSrc, flags);
  staging_buffer.Update(raw_data);
  stbi_image_free(raw_data);

  vk::ImageCreateInfo img_create_info;
  img_create_info.setExtent({ m_width, m_height, 1 });
  img_create_info.setFormat(format);
  img_create_info.setUsage(
      vk::ImageUsageFlagBits::eTransferSrc | vk::ImageUsageFlagBits::eTransferDst | vk::ImageUsageFlagBits::eSampled);
  img_create_info.setImageType(vk::ImageType::e2D);
  img_create_info.setMipLevels(m_mip_levels);
  img_create_info.setArrayLayers(1);
  img_create_info.setTiling(vk::ImageTiling::eOptimal);
  img_create_info.setSamples(vk::SampleCountFlagBits::e1);
  img_create_info.setSharingMode(vk::SharingMode::eExclusive);

  auto [img_result, image] = device->GetLogicalDeviceHandle().createImageUnique(img_create_info);
  assert(img_result == vk::Result::eSuccess);
  m_image = std::move(image);

  m_memory = CreateImageMemory(device, *m_image);

  TransitionImageLayout(*m_image, PRE_MIP_LAYOUT, m_mip_levels);

  auto& transfer_cmd_buffer = context.BeginSingleTimeCommands();

  vk::BufferImageCopy image_copy;
  image_copy.setBufferOffset(0);
  image_copy.setBufferRowLength(0);
  image_copy.setBufferImageHeight(0);
  image_copy.imageSubresource.aspectMask     = vk::ImageAspectFlagBits::eColor;
  image_copy.imageSubresource.baseArrayLayer = 0;
  image_copy.imageSubresource.layerCount     = 1;
  image_copy.imageSubresource.mipLevel       = 0;
  image_copy.setImageExtent({ m_width, m_height, 1 });
  image_copy.setImageOffset({ 0, 0, 0 });

  transfer_cmd_buffer.copyBufferToImage(
      staging_buffer.GetBuffer(), *m_image, vk::ImageLayout::eTransferDstOptimal, image_copy);
  context.EndSingleTimeCommands();

  GenerateMipMaps();

  m_image_view =
      CreateImageView(device, *m_image, format, vk::ImageAspectFlagBits::eColor, vk::ImageViewType::e2D, m_mip_levels);
  CreateDescriptorSet();
}

void Texture::CreateDescriptorSet()
{
  auto& context = fpr::Context::Get();
  auto  device  = context.GetDevice();

  vk::DescriptorSetLayout sampler_layout = context.m_render_graph->DescriptorSetLayouts.Get("Sampler");

  vk::DescriptorSetAllocateInfo desc_set_alloc_info;
  desc_set_alloc_info.descriptorPool = context.GetDescriptorPool();
  desc_set_alloc_info.setDescriptorSetCount(1);
  desc_set_alloc_info.setSetLayouts(sampler_layout);
  auto [alloc_result, descriptors] = device->GetLogicalDeviceHandle().allocateDescriptorSetsUnique(desc_set_alloc_info);

  m_descriptor_set = std::move(descriptors[0]);
  assert(alloc_result == vk::Result::eSuccess);

  vk::DescriptorImageInfo image_info;
  image_info.setImageLayout(vk::ImageLayout::eShaderReadOnlyOptimal);
  image_info.setImageView(*m_image_view);
  image_info.setSampler(context.m_render_graph->Samplers.Get("Sampler"));

  vk::WriteDescriptorSet write_set;
  write_set.setImageInfo(image_info);
  write_set.setDstSet(*m_descriptor_set);
  write_set.setDescriptorType(vk::DescriptorType::eCombinedImageSampler);
  device->GetLogicalDeviceHandle().updateDescriptorSets(write_set, nullptr);
}

void TransitionImageLayout(vk::Image image, const LayoutTransitionRules& transition_rules, uint32_t mip_levels)
{
  auto& context = fpr::Context::Get();

  auto& single_use_cmd_buffer = context.BeginSingleTimeCommands();

  vk::ImageMemoryBarrier img_mem_barrier;
  img_mem_barrier.setOldLayout(transition_rules.old_layout);
  img_mem_barrier.setNewLayout(transition_rules.new_layout);
  img_mem_barrier.setImage(image);

  img_mem_barrier.subresourceRange.aspectMask     = transition_rules.image_aspect;
  img_mem_barrier.subresourceRange.baseMipLevel   = 0;
  img_mem_barrier.subresourceRange.baseArrayLayer = 0;
  img_mem_barrier.subresourceRange.layerCount     = 1;
  img_mem_barrier.subresourceRange.levelCount     = mip_levels;

  img_mem_barrier.srcAccessMask = transition_rules.src_access;
  img_mem_barrier.dstAccessMask = transition_rules.dst_access;

  single_use_cmd_buffer.pipelineBarrier(
      transition_rules.src_stage,
      transition_rules.dst_stage,
      vk::DependencyFlagBits::eByRegion,
      nullptr,
      nullptr,
      img_mem_barrier);

  context.EndSingleTimeCommands();
}
} // namespace fpr