#include <GraphicsContext/GraphicsContext.hpp>
#include <Vk/CommandBuffer/CommandBuffer.hpp>
#include <Vk/DescriptorPool/DescriptorPool.hpp>
#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <Vk/ShaderConst.hpp>
#include <VoxelTree/VoxelTree.hpp>
#include <BlockRegistry/BlockRegistry.hpp>
#include <vulkan/vulkan_core.h>

void VoxelTree::markNodeForUpdate(NodeID nid)
{
    if (treeBuffer == nullptr)
    {
        // Nothing to do yet, because the tree is yet to be uploaded.
        return;
    }

    if (nodeUpdateIndexes.contains(nid))
    {
        updateList[nodeUpdateIndexes[nid]].node = nodeTable[nid];
        return;
    }

    uint32_t index = updateList.size();
    updateList.push({nid, nodeTable[nid]});
    nodeUpdateIndexes[nid] = index;
}

VoxelTree::VoxelTree(std::shared_ptr<Vk::Device> device) :
    device(device),
    updateList(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
{
    nodeTable = blockRegistry.getAtlasNodes();
    nodeTable.reserve(initialReservedNodes);
}

VoxelTree::NodeID VoxelTree::createNode(NodeID fill)
{
    if (!reclaimStack.empty())
    {
        NodeID result = reclaimStack.back();
        reclaimStack.resize(reclaimStack.size() - 1);
        nodeTable[result] = {fill, fill, fill, fill, fill, fill, fill, fill};
        markNodeForUpdate(result);
        return result;
    }

    NodeID result = nodeTable.size();
    nodeTable.push_back({fill, fill, fill, fill, fill, fill, fill, fill});
    markNodeForUpdate(result);
    return result;
}

void VoxelTree::releaseNode(NodeID nid)
{
    reclaimStack.push_back(nid);
}

VoxelTree::NodeID VoxelTree::getChildNode(NodeID parentNode, uint32_t index)
{
    return nodeTable[parentNode].childRefs[index];
}

void VoxelTree::setChildNode(NodeID parentNode, uint32_t index, NodeID childNode)
{
    Node &node = nodeTable[parentNode];
    if (node.childRefs[index] == childNode)
    {
        return;
    }

    node.childRefs[index] = childNode;
    markNodeForUpdate(parentNode);
}

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

    if (treeBuffer == nullptr)
    {
        std::shared_ptr<Vk::Buffer> stagingBuffer = std::make_shared<Vk::Buffer>(
            device,
            sizeof(Node) * nodeTable.size(),
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT
        );

        treeBuffer = std::make_shared<Vk::Buffer>(
            device,
            sizeof(Node) * initialReservedNodes,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
        );

        treeBufferCapacity = initialReservedNodes;

        memcpy(stagingBuffer->hostdata(), nodeTable.data(), sizeof(Node) * nodeTable.size());
        stagingBuffer->flush(0, sizeof(Node) * nodeTable.size());

        context.holdResource(stagingBuffer);

        cmdbuf.transitionBufferToTransferDst(*treeBuffer);
        cmdbuf.copyBuffer(*stagingBuffer, 0, *treeBuffer, 0, sizeof(Node) * nodeTable.size());
        cmdbuf.transitionBufferToShaderRead(*treeBuffer);

        descPool = std::make_shared<Vk::DescriptorPool>(
            device,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            2 * GraphicsContext::maxFramesInFlight,
            2 * GraphicsContext::maxFramesInFlight
        );

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

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

    if (nodeTable.size() >= treeBufferCapacity)
    {
        // We need to hold the previous descriptors in case they are still in use by a
        // rendering frame.
        context.holdResource(treeBuffer);
        context.holdResource(descSet);

        size_t prevCapacity = treeBufferCapacity;
        while (nodeTable.size() >= treeBufferCapacity)
        {
            if (treeBufferCapacity < slowAllocationThreshold)
            {
                treeBufferCapacity <<= 1;
            }
            else
            {
                treeBufferCapacity += slowAllocationStep;
            }
        }

        std::shared_ptr<Vk::Buffer> newTreeBuffer = std::make_shared<Vk::Buffer>(
            device,
            sizeof(Node) * treeBufferCapacity,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT
        );

        cmdbuf.barrier(
            *treeBuffer,
            VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            VK_ACCESS_SHADER_READ_BIT,
            VK_PIPELINE_STAGE_TRANSFER_BIT,
            VK_ACCESS_TRANSFER_READ_BIT
        );

        cmdbuf.copyBuffer(*treeBuffer, 0, *newTreeBuffer, 0, sizeof(Node) * prevCapacity);

        cmdbuf.barrier(
            *newTreeBuffer,
            VK_PIPELINE_STAGE_TRANSFER_BIT,
            VK_ACCESS_TRANSFER_WRITE_BIT,
            VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            VK_ACCESS_SHADER_READ_BIT
        );

        treeBuffer = newTreeBuffer;

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

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

    if (!updateList.empty())
    {
        uint32_t numUpdates = updateList.size();
        std::shared_ptr<Vk::Buffer> updateBuffer = updateList.unpack();
        nodeUpdateIndexes.clear();

        std::shared_ptr<Vk::DescriptorSet> updateDescSet = std::make_shared<Vk::DescriptorSet>(
            device,
            descPool,
            GraphicsContext::arrayDescriptorSetLayout
        );

        updateDescSet->setStorageBuffer(0, *updateBuffer, 0, VK_WHOLE_SIZE);

        context.holdResource(updateBuffer);
        context.holdResource(updateDescSet);

        cmdbuf.bindPipeline(VK_PIPELINE_BIND_POINT_COMPUTE, *GraphicsContext::svoUpdatePipeline);
        cmdbuf.bindDescriptorSet(
            VK_PIPELINE_BIND_POINT_COMPUTE,
            *GraphicsContext::svoUpdatePipeline,
            GLSL_SVOUPDATE_SVO_DESCRIPTOR_SET,
            *descSet
        );

        cmdbuf.bindDescriptorSet(
            VK_PIPELINE_BIND_POINT_COMPUTE,
            *GraphicsContext::svoUpdatePipeline,
            GLSL_SVOUPDATE_LIST_DESCRIPTOR_SET,
            *updateDescSet
        );

        cmdbuf.barrier(
            *treeBuffer,
            VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            VK_ACCESS_SHADER_READ_BIT,
            VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            VK_ACCESS_SHADER_WRITE_BIT
        );

        UpdatePushConst pushConst;
        pushConst.numUpdates = numUpdates;

        for (pushConst.updateBase = 0; pushConst.updateBase < pushConst.numUpdates; pushConst.updateBase += device->maxWorkGroupCount.x)
        {
            uint32_t updatesInIteration = std::min(device->maxWorkGroupCount.x, numUpdates - pushConst.updateBase);
            uint32_t numWorkGroups = (updatesInIteration + GLSL_SVOUPDATE_WORKGROUP_SIZE - 1) / GLSL_SVOUPDATE_WORKGROUP_SIZE;

            cmdbuf.pushConstants(
                *GraphicsContext::svoUpdatePipeline,
                VK_SHADER_STAGE_COMPUTE_BIT,
                0,
                sizeof(pushConst),
                &pushConst
            );

            cmdbuf.dispatch(glm::uvec3(numWorkGroups, 1, 1));
        }

        cmdbuf.barrier(
            *treeBuffer,
            VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            VK_ACCESS_SHADER_WRITE_BIT,
            VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            VK_ACCESS_SHADER_READ_BIT
        );
    }

    return *descSet;
}

size_t VoxelTree::getNodesInUse()
{
    return nodeTable.size() - reclaimStack.size();
}

VoxelTree::NodeID VoxelTree::releaseNodeIfUniform(NodeID nid)
{
    if (nid & GLSL_SVO_VOXEL_BIT)
    {
        return nid;
    }

    const Node &node = nodeTable[nid];
    if ((node.childRefs[0] & GLSL_SVO_VOXEL_BIT) == 0)
    {
        return nid;
    }

    for (size_t i = 1; i < 8; i++)
    {
        if (node.childRefs[i] != node.childRefs[0])
        {
            return nid;
        }
    }

    NodeID result = node.childRefs[0];
    releaseNode(nid);
    return result;
}