#ifndef GPU_MIRRORED_ARRAY_HPP_
#define GPU_MIRRORED_ARRAY_HPP_

#include <HostVisibleVector/HostVisibleVector.hpp>
#include <Vk/Buffer/Buffer.hpp>
#include <Vk/CommandBuffer/CommandBuffer.hpp>
#include <Vk/Device/Device.hpp>
#include <Vk/DescriptorPool/DescriptorPool.hpp>
#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <unordered_map>
#include <vector>
#include <Utilities/Extern.hpp>
#include <Settings/Settings.hpp>
#include <GraphicsContext/GraphicsContext.hpp>

/**
 * @brief A dynamic array which is mirrored by a device-local storage buffer.
 */
template<typename T>
class GPUMirroredArray final
{
private:
    static constexpr size_t slowAllocationThreshold = 1 << 25;
    static constexpr size_t slowAllocationStep = 1 << 24;

    size_t privSize = 0;
    size_t cpuCapacity = 0;
    size_t gpuCapacity = 0;

    std::shared_ptr<Vk::Device> device;
    HostVisibleVector<T> deltas;
    std::unique_ptr<Vk::Buffer> deviceLocalBuffer;
    std::shared_ptr<Vk::DescriptorPool> descPool;
    std::unique_ptr<Vk::DescriptorSet> descSet;
    std::vector<T> localVector;
    std::vector<VkBufferCopy> updateRegions;
    bool requiresFullFlush = true;

    // Maps an array index to the delta index.
    std::unordered_map<size_t, size_t> deltaIndexes;

public:
    GPUMirroredArray(std::shared_ptr<Vk::Device> device) :
        device(device),
        deltas(VK_BUFFER_USAGE_TRANSFER_SRC_BIT)
    {
        descPool = std::make_shared<Vk::DescriptorPool>(
            device,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            4,
            4
        );

        descSet = std::make_unique<Vk::DescriptorSet>(
            device,
            descPool,
            GraphicsContext::arrayDescriptorSetLayout
        );

        reserve(16);
    }

    const T* peek()
    {
        return localVector.data();
    }

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

    void reserve(size_t minCapacity)
    {
        if (cpuCapacity >= minCapacity)
        {
            return;
        }

        if (cpuCapacity == 0)
        {
            cpuCapacity = 1;
        }

        while (cpuCapacity < minCapacity)
        {
            if (cpuCapacity >= slowAllocationThreshold)
            {
                cpuCapacity += slowAllocationStep;
            }
            else
            {
                cpuCapacity <<= 1;
            }
        }

        localVector.reserve(cpuCapacity);
    }

    size_t size() const
    {
        return privSize;
    }

    size_t capacity() const
    {
        return cpuCapacity;
    }

    void replace(size_t index, T value)
    {
        localVector[index] = value;

        if (requiresFullFlush)
        {
            // Nothing more to do, we are already flushing the whole buffer.
            return;
        }

        if (!deltaIndexes.contains(index))
        {
            if (settings.validationLayers && deltas.size() > 1024)
            {
                requiresFullFlush = true;
                return;
            }

            if (deltas.size() > localVector.size()/8)
            {
                requiresFullFlush = true;
                return;
            }

            size_t deltaIndex = deltas.size();
            deltaIndexes[index] = deltaIndex;
            deltas.push(value);

            VkBufferCopy region{};
            region.srcOffset = sizeof(T) * deltaIndex;
            region.dstOffset = sizeof(T) * index;
            region.size = sizeof(T);

            updateRegions.push_back(region);
            return;
        }
        
        deltas[deltaIndexes[index]] = value;
    }

    void append(T value)
    {
        reserve(privSize + 1);
        size_t index = privSize++;
        localVector.push_back(value);
        replace(index, value);
    }

    void resize(size_t newSize)
    {
        reserve(newSize);
        localVector.resize(newSize);
        privSize = newSize;
    }

    Vk::DescriptorSet& getDescSet(Vk::CommandBuffer &cmdbuf)
    {
        auto &context = *GraphicsContext::instance;

        if (!requiresFullFlush && deltas.empty())
        {
            return *descSet;
        }

        if (requiresFullFlush)
        {
            requiresFullFlush = false;

            // Full flush, copy our data into a staging buffer and overwrite the device-local buffer completely.
            auto stagingBuffer = std::make_shared<Vk::Buffer>(
                device,
                sizeof(T) * privSize,
                VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
            );

            context.holdResource(stagingBuffer);

            std::memcpy(stagingBuffer->hostdata(), localVector.data(), sizeof(T) * privSize);

            if (cpuCapacity > gpuCapacity)
            {
                gpuCapacity = cpuCapacity;

                context.holdResource(std::move(deviceLocalBuffer));
                deviceLocalBuffer = std::make_unique<Vk::Buffer>(
                    device,
                    sizeof(T) * gpuCapacity,
                    VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
                    VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
                );

                context.holdResource(std::move(descSet));

                descSet = std::make_unique<Vk::DescriptorSet>(
                    device,
                    descPool,
                    GraphicsContext::arrayDescriptorSetLayout
                );

                descSet->setStorageBuffer(0, *deviceLocalBuffer, 0, VK_WHOLE_SIZE);
            }

            // Wait for everyone to finish work on the device-local buffer before we begin new work.
            cmdbuf.barrier(
                *deviceLocalBuffer,
                VK_PIPELINE_STAGE_TRANSFER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_VERTEX_SHADER_BIT,
                VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_TRANSFER_WRITE_BIT,
                VK_PIPELINE_STAGE_TRANSFER_BIT,
                VK_ACCESS_TRANSFER_WRITE_BIT
            );

            cmdbuf.copyBuffer(*stagingBuffer, 0, *deviceLocalBuffer, 0, sizeof(T) * privSize);

            // Finish the transfer before shaders start reading the device-local buffer.
            cmdbuf.barrier(
                *deviceLocalBuffer,
                VK_PIPELINE_STAGE_TRANSFER_BIT,
                VK_ACCESS_TRANSFER_WRITE_BIT,
                VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_VERTEX_SHADER_BIT,
                VK_ACCESS_SHADER_READ_BIT
            );
        }
        else if (!deltas.empty())
        {
            // We must wait for all work to complete on the old buffer before we read or write it.
            cmdbuf.barrier(
                *deviceLocalBuffer,
                VK_PIPELINE_STAGE_TRANSFER_BIT | VK_PIPELINE_STAGE_VERTEX_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                VK_ACCESS_TRANSFER_WRITE_BIT | VK_ACCESS_SHADER_READ_BIT,
                VK_PIPELINE_STAGE_TRANSFER_BIT,
                VK_ACCESS_TRANSFER_READ_BIT | VK_ACCESS_TRANSFER_WRITE_BIT
            );

            if (cpuCapacity > gpuCapacity)
            {
                auto oldBuffer = std::move(deviceLocalBuffer);
                deviceLocalBuffer = std::make_unique<Vk::Buffer>(
                    device,
                    sizeof(T) * cpuCapacity,
                    VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
                    VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
                );

                cmdbuf.copyBuffer(
                    *oldBuffer, 0,
                    *deviceLocalBuffer, 0,
                    sizeof(T) * gpuCapacity
                );

                context.holdResource(std::move(oldBuffer));

                context.holdResource(std::move(descSet));

                descSet = std::make_unique<Vk::DescriptorSet>(
                    device,
                    descPool,
                    GraphicsContext::arrayDescriptorSetLayout
                );

                descSet->setStorageBuffer(0, *deviceLocalBuffer, 0, VK_WHOLE_SIZE);

                gpuCapacity = cpuCapacity;

                // Wait for the above transfer to complete, before individual elements get patched.
                cmdbuf.barrier(
                    *deviceLocalBuffer,
                    VK_PIPELINE_STAGE_TRANSFER_BIT,
                    VK_ACCESS_TRANSFER_WRITE_BIT,
                    VK_PIPELINE_STAGE_TRANSFER_BIT,
                    VK_ACCESS_TRANSFER_WRITE_BIT
                );
            }

            auto deltaBuffer = deltas.unpack();
            context.holdResource(deltaBuffer);

            cmdbuf.copyBuffer(*deltaBuffer, *deviceLocalBuffer, updateRegions);

            // Wait for transfers to complete before shaders read the device-local buffer.
            cmdbuf.barrier(
                *deviceLocalBuffer,
                VK_PIPELINE_STAGE_TRANSFER_BIT,
                VK_ACCESS_TRANSFER_WRITE_BIT,
                VK_PIPELINE_STAGE_VERTEX_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                VK_ACCESS_SHADER_READ_BIT
            );
        }

        deltas.clear();
        deltaIndexes.clear();
        updateRegions.clear();

        return *descSet;
    }
};

#endif