#ifndef BLOCK_ATLAS_HPP_
#define BLOCK_ATLAS_HPP_

#include <Utilities/RGBABytes.hpp>
#include <Vk/Buffer/Buffer.hpp>
#include <Vk/DescriptorPool/DescriptorPool.hpp>
#include <vector>
#include <Material/Material.hpp>
#include <Vk/ShaderConst.hpp>
#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <VoxelTree/VoxelTree.hpp>
#include <BlockSprite/BlockSprite.hpp>

/**
 * @brief Builds and manages the block texture atlas.
 */
class BlockAtlas
{
public:
    typedef uint32_t SpriteID;

    struct Voxel
    {
        uint32_t color;
        uint32_t lightSpec;
        uint32_t bias;
        uint32_t refract;
        uint32_t frameDuration;
    };

    struct NodeSampleFace
    {
        uint32_t radiance = 0;
        uint32_t filter = 0;
    };

    struct NodeSampleFrame
    {
        NodeSampleFace faces[6];
    };

    struct NodeSample
    {
        NodeSampleFrame frames[8];
    };

    std::vector<VoxelTree::Node> atlasNodes;
    std::vector<NodeSample> nodeSamples;
    std::vector<Voxel> voxelPalette;
    std::shared_ptr<Vk::Buffer> paletteStagingBuffer;
    std::shared_ptr<Vk::Buffer> nodeSamplesStagingBuffer;
    std::shared_ptr<Vk::Buffer> voxelPaletteBuffer;
    std::shared_ptr<Vk::Buffer> nodeSamplesBuffer;
    std::shared_ptr<Vk::DescriptorPool> voxelPaletteDescPool;
    std::shared_ptr<Vk::DescriptorSet> voxelPaletteDescSet;

private:
    static constexpr uint32_t numNodeBuckets = 1 << 16;
    static constexpr uint32_t numVoxelBuckets = 1 << 16;

    std::vector<std::vector<uint32_t>> nodeHashSet;
    std::vector<std::vector<uint32_t>> voxelHashSet;
    size_t finalVoxelCount = 0;

    std::optional<VoxelTree::NodeID> tryCollapseNode(const VoxelTree::Node &node);
    NodeSample getNodeSample(VoxelTree::NodeID nodeId);
    VoxelTree::NodeID addNode(const VoxelTree::Node &node);
    RGBABytes decodeColor(uint32_t rgba);
    uint32_t encodeColor(VoxelSprite::Color &color);
    uint32_t encodeVoxelFrames(const std::vector<Voxel> &voxels);
    VoxelTree::NodeID makeVoxelNode(BlockSprite &sprite, glm::uvec3 pos);
    VoxelTree::NodeID makeSpriteNode(BlockSprite &sprite, VoxelTree::Order order, glm::uvec3 offset);
    bool areVoxelsSame(uint32_t nodeRef, const std::vector<Voxel> &voxels);

public:
    BlockAtlas();

    /**
     * @brief Add the specified sprite to the atlas and returns its ID.
     * 
     * Note that the sprite passed in will be modified by this method, for the purposes
     * of optimising the sprite.
     */
    SpriteID addSprite(BlockSprite &sprite);

    /**
     * @brief Get a const reference to the atlas nodes.
     */
    const std::vector<VoxelTree::Node>& getAtlasNodes();

    /**
     * @brief Returns the number of voxels in the palette.
     */
    size_t getNumVoxels();

    void loadVoxelPalette(std::atomic<unsigned int> &tasksDone);
    Vk::DescriptorSet& getVoxelPaletteDescSet();
};

#endif