#ifndef VOXEL_TREE_HPP_
#define VOXEL_TREE_HPP_

#include <GraphicsContext/GraphicsContext.hpp>
#include <Utilities/Extern.hpp>
#include <Vk/DescriptorPool/DescriptorPool.hpp>
#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <cstdint>
#include <vector>
#include <Vk/ShaderConst.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>
#include <HostVisibleVector/HostVisibleVector.hpp>

/**
 * @brief Represents the Sparse Voxel Octree (SVO) and allows its manipulation on the GPU.
 */
class VoxelTree
{
public:
    typedef uint32_t NodeID;
    typedef unsigned int Order;

    struct Node
    {
        NodeID childRefs[8];
    };

    struct Update
    {
        NodeID index;
        Node node;
    };

    struct UpdatePushConst
    {
        uint32_t updateBase;
        uint32_t numUpdates;
    };

private:
    static constexpr size_t slowAllocationThreshold = 1 << 25;
    static constexpr size_t slowAllocationStep = 1 << 24;

    std::shared_ptr<Vk::Device> device;
    std::vector<NodeID> reclaimStack;
    std::vector<Node> nodeTable;
    HostVisibleVector<Update> updateList;
    std::shared_ptr<Vk::Buffer> treeBuffer;
    size_t treeBufferCapacity;
    std::shared_ptr<Vk::DescriptorPool> descPool;
    std::shared_ptr<Vk::DescriptorSet> descSet;
    std::unordered_map<NodeID, uint32_t> nodeUpdateIndexes;

    void markNodeForUpdate(NodeID nid);

public:
    static constexpr Order rootOrder = GLSL_SVO_MAX_ORDER;
    static constexpr NodeID rootNode = 0;
    static constexpr size_t initialReservedNodes = 1 << 18;

    /**
     * @brief Returns the index of the child node containing coords.
     * 
     * @param coords The coordinates that the child nodes would contain.
     * @param order The order of the parent node (must be at least 1).
     */
    static constexpr uint32_t getChildNodeIndex(glm::uvec3 coords, Order order)
    {
        uint32_t xpart = (coords.x >> (order - 1)) & 1;
        uint32_t ypart = (coords.y >> (order - 1)) & 1;
        uint32_t zpart = (coords.z >> (order - 1)) & 1;

        return (zpart << 2) | (ypart << 1) | xpart;
    };

    VoxelTree(std::shared_ptr<Vk::Device> device);

    /**
     * @brief Allocates a new, empty node, and returns its ID.
     */
    NodeID createNode(NodeID fill);

    /**
     * @brief Release the specified NodeID, so that it may be recycled for a future node.
     */
    void releaseNode(NodeID nid);

    /**
     * @brief Return the ID of the specified child of the node.
     *
     * Returns `rootNode` if the child does not exist.
     */
    NodeID getChildNode(NodeID parentNode, uint32_t index);

    /**
     * @brief Sets the ID of the specified child node.
     * 
     * @param parentNode The parent node.
     * @param index The child index.
     * @param childNode The child node ID; if `rootNode`, then the child is removed.
     */
    void setChildNode(NodeID parentNode, uint32_t index, NodeID childNode);

    /**
     * @brief Record commands to update the tree and then return its descriptor set.
     */
    Vk::DescriptorSet& getDescSet(Vk::CommandBuffer &cmdbuf);

    /**
     * @brief Release the specified node if it is uniform, and return the contained voxel.
     */
    NodeID releaseNodeIfUniform(NodeID nid);

    size_t getNodesInUse();
};

#endif