#ifndef BLOCK_REGISTRY_HPP_
#define BLOCK_REGISTRY_HPP_

#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <unordered_map>
#include <vector>
#include <functional>

#include <Block/Block.hpp>
#include <BlockAtlas/BlockAtlas.hpp>
#include <World/World.hpp>
#include <Block/SpriteBlock/Stairs/Stairs.hpp>

class Slab;
class SpriteBlock;

/**
 * @brief Registry of blocks; maps block names to block implementations.
 */
class BlockRegistry
{
private:
    // Flags for registerBlock().
    static constexpr int withSlab = 1 << 0;
    static constexpr int withStairs = 1 << 1;

    std::unordered_map<std::string, Block*> blocks;
    std::vector<Block*> orderedBlocks;
    Block *blockAir;
    BlockAtlas atlas;

    void registerDyedBlocks(std::string dyeColorName);

    /**
     * @brief Register a block.
     * 
     * @tparam TBlock The block class.
     * @param name The block name.
     * @param args Arguments to the constructor (besides the name).
     */
    template<typename TBlock, typename... TArgs>
    TBlock& registerBlock(int flags, std::string name, TArgs&& ...args)
    {
        TBlock *block = new TBlock(name, args...);
        blocks[name] = block;
        orderedBlocks.push_back(block);

        if (flags & withSlab)
        {
            registerBlock<Slab>(name + "Slab", dynamic_cast<SpriteBlock&>(*block));
        }

        if (flags & withStairs)
        {
            registerBlock<Stairs>(name + "Stairs", dynamic_cast<SpriteBlock&>(*block));
        }

        return *block;
    }

    template<typename TBlock, typename... TArgs>
    TBlock& registerBlock(std::string name, TArgs&& ...args)
    {
        return registerBlock<TBlock>(0, name, args...);
    }

    void registerDecorativeUniformStone(std::string name, std::string displayName, float hardness, std::string materialName);
    void registerDecorativeUniformGlass(std::string name, std::string displayName, float hardness, std::string materialName);

public:
    ~BlockRegistry();

    /**
     * @brief Perform the registration of all blocks.
     */
    void registerAllBlocks();

    /**
     * @brief Return a reference to the air block.
     * 
     * @return The air block.
     */
    Block& getAirBlock();

    /**
     * @brief Get a block given a name.
     * 
     * @param name The block name.
     * @return Reference to the block implementation.
     */
    Block& getBlockByName(std::string name);

    /**
     * @brief Add the callbacks to be called during preloading.
     */
    void addLoadingTasks(std::vector<std::function<void(std::atomic<unsigned int> &)>> &tasks);

    /**
     * @brief Add block names to the vector.
     */
    void assignBlockIds(std::vector<std::string> &ids);

    /**
     * @brief Returns true if the specified block name is valid.
     */
    bool isValidBlock(std::string name);

    /**
     * @brief Get stacks for every block.
     * 
     * @param stacks Stacks will be added here.
     */
    void getAllStacks(std::vector<ItemStack> &stacks);

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

    /**
     * @brief Call the callback on each block, in a deterministic order.
     */
    void forEachBlock(std::function<void(Block&)> callback);

    void loadBoxesOnServer();
    int getTotalBlocks();
    size_t getNumVoxels();
    Vk::DescriptorSet& getVoxelPaletteDescSet();
};

extern BlockRegistry blockRegistry;

#endif