#include "XCEngine/RHI/Vulkan/VulkanBuffer.h"

#include "XCEngine/RHI/Vulkan/VulkanDevice.h"

#include <cstring>

namespace XCEngine {
namespace RHI {

VulkanBuffer::~VulkanBuffer() {
    Shutdown();
}

bool VulkanBuffer::Initialize(VulkanDevice* device, const BufferDesc& desc, VkBufferUsageFlags usageFlags, VkMemoryPropertyFlags memoryProperties) {
    if (device == nullptr || device->GetDevice() == VK_NULL_HANDLE || desc.size == 0) {
        return false;
    }

    m_deviceOwner = device;
    m_device = device->GetDevice();
    m_size = desc.size;
    m_stride = desc.stride;
    m_bufferType = static_cast<BufferType>(desc.bufferType);

    VkBufferCreateInfo bufferInfo = {};
    bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferInfo.size = static_cast<VkDeviceSize>(desc.size);
    bufferInfo.usage = usageFlags;
    bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

    if (vkCreateBuffer(m_device, &bufferInfo, nullptr, &m_buffer) != VK_SUCCESS) {
        return false;
    }

    VkMemoryRequirements memoryRequirements = {};
    vkGetBufferMemoryRequirements(m_device, m_buffer, &memoryRequirements);

    VkMemoryAllocateInfo allocateInfo = {};
    allocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocateInfo.allocationSize = memoryRequirements.size;
    allocateInfo.memoryTypeIndex = device->FindMemoryType(memoryRequirements.memoryTypeBits, memoryProperties);
    if (allocateInfo.memoryTypeIndex == UINT32_MAX) {
        Shutdown();
        return false;
    }

    if (vkAllocateMemory(m_device, &allocateInfo, nullptr, &m_memory) != VK_SUCCESS) {
        Shutdown();
        return false;
    }

    if (vkBindBufferMemory(m_device, m_buffer, m_memory, 0) != VK_SUCCESS) {
        Shutdown();
        return false;
    }

    return true;
}

void VulkanBuffer::Shutdown() {
    if (m_mappedData != nullptr && m_device != VK_NULL_HANDLE && m_memory != VK_NULL_HANDLE) {
        vkUnmapMemory(m_device, m_memory);
        m_mappedData = nullptr;
    }

    if (m_buffer != VK_NULL_HANDLE && m_device != VK_NULL_HANDLE) {
        vkDestroyBuffer(m_device, m_buffer, nullptr);
        m_buffer = VK_NULL_HANDLE;
    }

    if (m_memory != VK_NULL_HANDLE && m_device != VK_NULL_HANDLE) {
        vkFreeMemory(m_device, m_memory, nullptr);
        m_memory = VK_NULL_HANDLE;
    }

    m_deviceOwner = nullptr;
    m_device = VK_NULL_HANDLE;
    m_size = 0;
    m_stride = 0;
    m_state = ResourceStates::Common;
}

void* VulkanBuffer::Map() {
    if (m_device == VK_NULL_HANDLE || m_memory == VK_NULL_HANDLE) {
        return nullptr;
    }

    if (m_mappedData == nullptr) {
        if (vkMapMemory(m_device, m_memory, 0, VK_WHOLE_SIZE, 0, &m_mappedData) != VK_SUCCESS) {
            m_mappedData = nullptr;
        }
    }

    return m_mappedData;
}

void VulkanBuffer::Unmap() {
    if (m_mappedData != nullptr && m_device != VK_NULL_HANDLE && m_memory != VK_NULL_HANDLE) {
        vkUnmapMemory(m_device, m_memory);
        m_mappedData = nullptr;
    }
}

void VulkanBuffer::SetData(const void* data, size_t size, size_t offset) {
    if (data == nullptr || size == 0 || offset + size > m_size) {
        return;
    }

    const bool wasMapped = m_mappedData != nullptr;
    void* mappedData = Map();
    if (mappedData == nullptr) {
        return;
    }

    std::memcpy(static_cast<uint8_t*>(mappedData) + offset, data, size);

    if (!wasMapped) {
        Unmap();
    }
}

} // namespace RHI
} // namespace XCEngine