#include <stdafx.h>
#include "Buffer/Buffer.h"
#include "Context.h"


namespace fpr
{
Buffer::Buffer(vk::DeviceSize buffer_size, vk::BufferUsageFlags usage_flags, vk::MemoryPropertyFlags mem_flags):
    m_buffer_size(buffer_size), m_usage(usage_flags), m_memory_flags(mem_flags)
{
  auto&                context = fpr::Context::Get();
  auto&                device  = context.GetDevice()->GetLogicalDeviceHandle();
  vk::BufferCreateInfo buffer_create_info{};
  buffer_create_info.size        = buffer_size;
  buffer_create_info.usage       = usage_flags;
  buffer_create_info.sharingMode = vk::SharingMode::eExclusive;

  auto [buffer_result, buffer] = device.createBufferUnique(buffer_create_info);

  assert(("Failed to create buffer.", buffer_result == vk::Result::eSuccess));

  vk::MemoryRequirements mem_req = device.getBufferMemoryRequirements(*buffer);
  vk::MemoryAllocateInfo mem_alloc_info{};
  const uint32_t         mem_type_index =
      FindMemoryTypeIndex(context.GetDevice()->GetPhysicalDeviceHandle(), mem_req.memoryTypeBits, mem_flags);

  [[maybe_unused]] bool found = mem_type_index != std::numeric_limits<uint32_t>::max();
  assert(("Memory type index not found.", found));

  mem_alloc_info.allocationSize = mem_req.size;

  mem_alloc_info.memoryTypeIndex = mem_type_index;

  auto [alloc_result, buffer_memory] = device.allocateMemoryUnique(mem_alloc_info);
  assert(("Failed to allocate buffer memory.", alloc_result == vk::Result::eSuccess));
  m_buffer = std::move(buffer);
  m_memory = std::move(buffer_memory);

  [[maybe_unused]] vk::Result bind_result = device.bindBufferMemory(*m_buffer, *m_memory, 0);

  assert(("Failed to bind buffer memory.", bind_result == vk::Result::eSuccess));
}

Buffer::~Buffer()
{
  //Most buffers are constructed and permanently mapped. Therefore, check if they are still mapped upon destructiion.
  if(mapped_data)
    Flush();
}

vk::DeviceSize Buffer::GetSize()
{
  return m_buffer_size;
}

vk::Buffer& Buffer::GetBuffer()
{
  return *m_buffer;
}

uint32_t FindMemoryTypeIndex(
    vk::PhysicalDevice      physical_device,
    uint32_t                type_field,
    vk::MemoryPropertyFlags type_flags)
{
  vk::PhysicalDeviceMemoryProperties device_mem_properties = physical_device.getMemoryProperties();

  for(uint32_t index = 0; index < device_mem_properties.memoryTypeCount; ++index)
  {
    bool memory_type_index_found = IsNthBitSet(type_field, index) &&
        IsOnlyNthBitSet(static_cast<uint32_t>(device_mem_properties.memoryTypes[index].propertyFlags),
                        static_cast<uint32_t>(type_flags));
    if(memory_type_index_found)
      return index;
  }
  return std::numeric_limits<uint32_t>::max();
}

bool IsNthBitSet(uint32_t bit_field, uint32_t nth_bit) FPR_NOEXCEPT
{
  return (bit_field & (1 << nth_bit));
};

bool IsOnlyNthBitSet(uint32_t bit_field, uint32_t nth_bit) FPR_NOEXCEPT
{
  return (bit_field & nth_bit) == nth_bit;
}


vk::UniqueDeviceMemory CreateImageMemory(fpr::Device* device, vk::Image image)
{
  vk::MemoryRequirements mem_req = device->GetLogicalDeviceHandle().getImageMemoryRequirements(image);

  vk::MemoryAllocateInfo mem_alloc_info{};
  mem_alloc_info.allocationSize  = mem_req.size;
  mem_alloc_info.memoryTypeIndex = fpr::FindMemoryTypeIndex(
      device->GetPhysicalDeviceHandle(), mem_req.memoryTypeBits, vk::MemoryPropertyFlagBits::eDeviceLocal);
  auto [alloc_result, img_mem] = device->GetLogicalDeviceHandle().allocateMemoryUnique(mem_alloc_info);
  assert(("Failed to allocate image memory!", alloc_result == vk::Result::eSuccess));
  [[maybe_unused]] vk::Result bind_result = device->GetLogicalDeviceHandle().bindImageMemory(image, *img_mem, 0);

  assert(("Failed to bind image memory!", bind_result == vk::Result::eSuccess));

  return std::move(img_mem);
}

void Buffer::Transfer(Buffer* other)
{
  vk::CommandBuffer& transfer_cmd_buff = fpr::Context::Get().BeginSingleTimeCommands();

  vk::BufferCopy copy_info;
  copy_info.size = GetSize();
  transfer_cmd_buff.copyBuffer(GetBuffer(), other->GetBuffer(), copy_info);
  fpr::Context::Get().EndSingleTimeCommands();
}

void Buffer::Update(const void* new_data)
{
  void* mapped = GetMapped();
  memcpy(mapped, new_data, m_buffer_size);

  Flush();
}

void Buffer::Flush()
{
  vk::MappedMemoryRange range;
  range.memory            = *m_memory;
  range.size              = m_buffer_size;
  [[maybe_unused]] vk::Result flush_result = fpr::Context::Get().GetDevice()->GetLogicalDeviceHandle().flushMappedMemoryRanges(range);
  
  assert(("Failed to flush mapped memory.",flush_result == vk::Result::eSuccess));
}

void* Buffer::GetMapped()
{
  //Get mapped memory. If not mapped, map it and return pointer to memory.
  if(!mapped_data) [[unlikely]]
  {
    vk::Result map_result;

    std::tie(map_result, mapped_data) =
        fpr::Context::Get().GetDevice()->GetLogicalDeviceHandle().mapMemory(*m_memory, 0, m_buffer_size);

    assert(("Failed to map memory.", map_result == vk::Result::eSuccess));
  }

  return mapped_data;
}

} // namespace fpr