#include <stdafx.h>
#include "Buffer/Buffer.h"
#include "Context.h"
#include "Memory.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;
}

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