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

#include "buffer.h"

#include "device.h"

void
Buffer::Destroy(const Device *device)
{
    if (!IsValid() || !IsOwned())
        return;

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

void
Buffer::Allocate(const Device *device, usize size, vk::BufferUsageFlags bufferUsage,
                 VmaAllocationCreateFlags allocationFlags, VmaMemoryUsage memoryUsage, cstr name)
{
    assert(!IsValid());
    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);

    vk::MemoryPropertyFlags memoryPropertyFlags;
    vmaGetAllocationMemoryProperties(device->m_Allocator, allocation,
                                     Recast<VkMemoryPropertyFlags *>(&memoryPropertyFlags));
    bool hostAccessible = Cast<bool>(memoryPropertyFlags & vk::MemoryPropertyFlagBits::eHostVisible);

    m_Buffer = buffer;
    m_Size_ = size | VALID_BUFFER_BIT | OWNED_BIT;
    m_Allocation = allocation;
    m_Mapped = Cast<u8 *>(allocationInfo.pMappedData);

    device->SetName(m_Buffer, name);
}

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

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

            vmaUnmapMemory(device->m_Allocator, m_Allocation);
            m_Mapped = nullptr;
        }
    }
    else
    {
        memcpy(m_Mapped + offset, 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
StorageBuffer::Init(const Device *device, usize size, bool hostVisible, cstr name)
{
    Init(device, size, hostVisible, false, name);
}

void
StorageBuffer::Init(const Device *device, usize size, bool hostVisible, bool deviceAddress, cstr name)
{
    vk::BufferUsageFlags usage = vk::BufferUsageFlagBits::eStorageBuffer;
    if (deviceAddress)
    {
        usage |= vk::BufferUsageFlagBits::eShaderDeviceAddress;   
    }
    if (hostVisible)
    {
        Allocate(device, size, usage,
                 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);
    }
    else
    {
        usage |= vk::BufferUsageFlagBits::eTransferDst;
        Allocate(device, size, usage, 0,
                 VMA_MEMORY_USAGE_AUTO, name);
    }
}

void
StorageIndexBuffer::Init(const Device *device, usize size, bool hostVisible, bool deviceAddress, cstr name)
{
    vk::BufferUsageFlags usage = vk::BufferUsageFlagBits::eStorageBuffer | vk::BufferUsageFlagBits::eIndexBuffer;
    if (deviceAddress)
    {
        usage |= vk::BufferUsageFlagBits::eShaderDeviceAddress;
    }
    if (hostVisible)
    {
        Allocate(device, size, usage,
                 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);
    }
    else
    {
        usage |= vk::BufferUsageFlagBits::eTransferDst;
        Allocate(device, size, usage, 0, VMA_MEMORY_USAGE_AUTO, name);
    }
}

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

void
IndexBuffer::Init(const Device *device, usize size, cstr name)
{
    Allocate(device, size, vk::BufferUsageFlagBits::eIndexBuffer | vk::BufferUsageFlagBits::eTransferDst,
             0, 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);
}