#pragma once

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

#include "GPU_GarbageCollector.h"
#include "Delegate/Delegate.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;
    
    VkBuffer create_buffer(VkDeviceSize size) const
    {
        return DeviceFunctions::create_buffer(
            device_.get_native(),
            size,
            VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
            physical_device_.get_unique_family_indices()
        );
    }
    VkDeviceMemory allocate_device_memory(VkBuffer buffer) const
    {
        return DeviceFunctions::allocate_device_memory(
            physical_device_.get_native(),
            device_.get_native(),
            buffer,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
        );
    }
public:
    struct BufferDataDelegate
    {
        VkBuffer buffer;
        VkDeviceMemory device_memory;
        VkDeviceSize size;
    };
    // Параметр - дескрипторы нового буфера
    Delegate<const BufferDataDelegate&> OnResize;
    
    explicit UniformBuffer( const VkObjects::Device& device,
                            const VkObjects::PhysicalDevice& physical_device,
                            const std::shared_ptr<GPU_GarbageCollector>& gc,
                            size_t count = 1):
    device_(device), physical_device_(physical_device), count_(count), gc_(gc)
    {
        if (count == 0)
            throw std::runtime_error("Count of 0 is not allowed to be 0");
        
        const VkDeviceSize size = sizeof(T) * count;
        buffer_ = create_buffer(size);
        device_memory_ = allocate_device_memory(buffer_);
        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!");
    }
    
    void resize(size_t new_count)
    {
        if (new_count == count_)
            return;
        
        if (new_count == 0)
            throw std::runtime_error("Count of 0 is not allowed to be 0");
        
        const VkDeviceSize new_size = sizeof(T) * new_count;
        void* new_map = nullptr;
        
        VkBuffer new_buffer = create_buffer(new_size);
        VkDeviceMemory new_memory = allocate_device_memory(new_buffer);
        vkBindBufferMemory(device_.get_native(), new_buffer, new_memory, 0);
        vkMapMemory(device_.get_native(), new_memory, 0, new_size, 0, &new_map);
        
        memcpy_s(new_map, new_size, map_, sizeof(T) * count_);
        
        vkUnmapMemory(device_.get_native(), device_memory_);
        gc_->AddGarbage({ buffer_, device_memory_ });
        BufferDataDelegate data{ new_buffer, new_memory, new_size };
        OnResize.call(data);
        
        buffer_ = new_buffer;
        device_memory_ = new_memory;
        map_ = new_map;
        count_ = new_count;
    }
};