#pragma once

#include <memory>
#include <stdexcept>
#include <vulkan/vulkan_core.h>

#include "GPU_GarbageCollector.h"
#include "DeviceFunctions/DeviceFunctions.hpp"
#include "VkObjects/Device.hpp"
#include "VkObjects/PhysicalDevice.hpp"


template <typename T>
class UniformBuffer
{
    const VkObjects::Device& device_;
    const VkObjects::PhysicalDevice& physical_device_;
    size_t count_;
    std::shared_ptr<GPU_GarbageCollector> gc_;
    
    VkBuffer buffer_ = nullptr;
    VkDeviceMemory device_memory_ = nullptr;
    void* map_ = nullptr;
public:
    explicit UniformBuffer( const VkObjects::Device& device,
                            const VkObjects::PhysicalDevice& physical_device,
                            const std::shared_ptr<GPU_GarbageCollector>& gc,
                            size_t count):
    device_(device), physical_device_(physical_device), count_(count), gc_(gc)
    {
        const VkDeviceSize size = sizeof(T) * count;
        buffer_ = DeviceFunctions::create_buffer(
            device.get_native(),
            size,
            VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
            physical_device.get_unique_family_indices()
        );
        device_memory_ = DeviceFunctions::allocate_device_memory(
            physical_device.get_native(),
            device.get_native(),
            buffer_,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
        );
        vkBindBufferMemory(device.get_native(), buffer_, device_memory_, 0);
    
        vkMapMemory(device.get_native(), device_memory_, 0, size, 0, &map_);
    }
    
    virtual ~UniformBuffer()
    {
        if (device_memory_ != nullptr && map_ != nullptr)
        {
            vkUnmapMemory(device_.get_native(), device_memory_);
            map_ = nullptr;
        }
    
        gc_->AddGarbage({buffer_, device_memory_});
    }
    
    operator T*() const
    {
        return map_;
    }
    
    operator void*() const
    {
        return map_;
    }
    
    operator VkBuffer() const
    {
        return buffer_;
    }
    
    T& operator[](size_t index)
    {
        if (index < count_)
            return map_[index];
        throw std::runtime_error("Index out of range!");
    }
};