#ifndef HOST_VISIBLE_VECTOR_HPP_
#define HOST_VISIBLE_VECTOR_HPP_

#include <Vk/Buffer/Buffer.hpp>
#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <memory>

/**
 * @brief A vector which allocates Vulkan host-visible memory.
 *
 * Once you fill the vector with data, you can call `unpack()` to get the buffer. A new
 * buffer with the same capacity will then be allocated for the vector, and the vector
 * itself cleared.
 */
template<typename T>
class HostVisibleVector final
{
private:
    static constexpr size_t initialCapacity = 128;

    std::shared_ptr<Vk::Device> device;
    std::shared_ptr<Vk::Buffer> buffer;
    VkBufferUsageFlags usageFlags;

    size_t privCapacity = initialCapacity;
    size_t privSize = 0;

    void grow()
    {
        privCapacity <<= 1;

        std::shared_ptr<Vk::Buffer> newBuffer = std::make_shared<Vk::Buffer>(
            device,
            sizeof(T) * privCapacity,
            usageFlags,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT
        );

        memcpy(newBuffer->hostdata(), buffer->hostdata(), sizeof(T) * privSize);

        buffer = newBuffer;
    }

public:
    HostVisibleVector(VkBufferUsageFlags usageFlags) : usageFlags(usageFlags)
    {
        device = GraphicsContext::instance->getDevice();

        buffer = std::make_shared<Vk::Buffer>(
            device,
            sizeof(T) * privCapacity,
            usageFlags,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT
        );
    }

    HostVisibleVector(const HostVisibleVector<T> &) = delete;

    T* data()
    {
        return (T*) buffer->hostdata();
    }

    T& operator[](size_t index)
    {
        return data()[index];
    }

    size_t capacity() const
    {
        return privCapacity;
    }

    size_t size() const
    {
        return privSize;
    }

    void push(T newElement)
    {
        while (privSize >= privCapacity)
        {
            grow();
        }

        data()[privSize++] = newElement;
    }

    void clear()
    {
        privSize = 0;
    }

    std::shared_ptr<Vk::Buffer> unpack()
    {
        std::shared_ptr<Vk::Buffer> oldBuffer = buffer;

        buffer = std::make_shared<Vk::Buffer>(
            device,
            sizeof(T) * privCapacity,
            usageFlags,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT
        );

        oldBuffer->flush(0, sizeof(T) * privSize);

        privSize = 0;

        return oldBuffer;
    }

    bool empty()
    {
        return privSize == 0;
    }
};

#endif