// =============================================
//  Aster: buffer.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "buffer.h"

#include "device.h"
#include <ranges>

void
Buffer::Destroy(const Device *device)
{
    if (!m_Buffer)
        return;

    vmaDestroyBuffer(device->m_Allocator, m_Buffer, m_Allocation);
    m_Buffer = nullptr;
}

void
Buffer::Allocate(const Device *device, usize size, vk::BufferUsageFlags bufferUsage,
                 VmaAllocationCreateFlags allocationFlags, VmaMemoryUsage memoryUsage, cstr name)
{
    assert(size <= SIZE_MASK);

    vk::BufferCreateInfo bufferCreateInfo = {
        .size = size,
        .usage = bufferUsage,
        .sharingMode = vk::SharingMode::eExclusive,
    };
    const VmaAllocationCreateInfo allocationCreateInfo = {
        .flags = allocationFlags,
        .usage = memoryUsage,
    };

    VkBuffer buffer;
    VmaAllocation allocation;
    VmaAllocationInfo allocationInfo;
    auto result = Cast<vk::Result>(vmaCreateBuffer(device->m_Allocator, Recast<VkBufferCreateInfo *>(&bufferCreateInfo),
                                                   &allocationCreateInfo, &buffer, &allocation, &allocationInfo));
    ERROR_IF(Failed(result), "Could not allocate buffer. Cause: {}", result) THEN_ABORT(result);

    m_Buffer = buffer;
    m_Size = size & SIZE_MASK;
    m_Allocation = allocation;
    m_Mapped = allocationInfo.pMappedData;

    vk::MemoryPropertyFlags memoryPropertyFlags;
    vmaGetAllocationMemoryProperties(device->m_Allocator, allocation,
                                     Recast<VkMemoryPropertyFlags *>(&memoryPropertyFlags));
    if (memoryPropertyFlags & vk::MemoryPropertyFlagBits::eHostVisible)
    {
        m_Size |= HOST_ACCESSIBLE_BIT;
    }

    device->SetName(m_Buffer, name);
}

void
Buffer::Write(const Device *device, usize offset, usize size, void *data)
{
    assert(IsHostVisible());

    if (!IsMapped())
    {
        auto result = Cast<vk::Result>(vmaMapMemory(device->m_Allocator, m_Allocation, &m_Mapped));
        ERROR_IF(Failed(result), "Memory mapping failed. Cause: {}", result);
        if (!Failed(result))
        {
            memcpy(m_Mapped, data, size);

            vmaUnmapMemory(device->m_Allocator, m_Allocation);
            m_Mapped = nullptr;
        }
    }
    else
    {
        memcpy(m_Mapped, data, size);
    }

    // TODO: Debug this.
    // auto result = Cast<vk::Result>(vmaCopyMemoryToAllocation(device->m_Allocator, &data, m_Allocation, 0, size));
    // ERROR_IF(Failed(result), "Writing to buffer failed. Cause: {}", result) THEN_ABORT(result);
}

void
UniformBuffer::Init(const Device *device, const usize size, const cstr name)
{
    Allocate(device, size, vk::BufferUsageFlagBits::eUniformBuffer,
             VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT |
                 VMA_ALLOCATION_CREATE_HOST_ACCESS_ALLOW_TRANSFER_INSTEAD_BIT | VMA_ALLOCATION_CREATE_MAPPED_BIT,
             VMA_MEMORY_USAGE_AUTO, name);
}

void
VertexBuffer::Init(const Device *device, usize size, cstr name)
{
    Allocate(device, size, vk::BufferUsageFlagBits::eVertexBuffer | vk::BufferUsageFlagBits::eTransferDst,
             VMA_ALLOCATION_CREATE_DEDICATED_MEMORY_BIT, VMA_MEMORY_USAGE_AUTO, name);
}

void
StagingBuffer::Init(const Device *device, usize size, cstr name)
{
    Allocate(device, size, vk::BufferUsageFlagBits::eTransferSrc,
             VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT | VMA_ALLOCATION_CREATE_MAPPED_BIT,
             VMA_MEMORY_USAGE_AUTO, name);
}