FeatherMC/src/nbt/NBT.h

#pragma once

#include <stddef.h>

#include <iostream>
#include <vector>
#include <iterator>
#include <type_traits>

#include "NBTInternal.h"

struct nbt_node;
struct list_head;

namespace NBT
{   
    // NBT-specific data buffer.
    // TODO: could replace with some project-wide utility class
    struct DataBuffer
    {
        uint8_t* data;
        size_t length;
    
        ~DataBuffer();
    };

    // Represents a generic NBT tag of any type
    class Tag
    {
    protected:
        nbt_node* m_node;
        Tag(nbt_node *node) : m_node(node) {}

        static Tag CreateTag(nbt_node* node);
        static Tag CreateListTag(nbt_node* node);
    public:

        // This will be "<null>" if the tag is nameless
        const char* GetName() const;

        // Get the number of iterable immediate children
        size_t GetLength() const;

        // Get uncompressed NBT binary data
        const DataBuffer GetData() const;

        virtual const char* GetValueString() const;

        const char* ToString() const;

        operator bool() const;
    };

    // Operator overload for ostream
    std::ostream& operator<<(std::ostream& stream, const Tag& tag);
    
    template <typename T>
    class DataTag : public Tag
    {
    public:
        DataTag<T>(nbt_node* node) : Tag(node) {}

        T GetValue() const;

        inline operator T() const { return GetValue(); }
    };

    // Standard value types
    using ByteTag = DataTag<int8_t>;
    using ShortTag = DataTag<int16_t>;
    using IntTag = DataTag<int32_t>;
    using LongTag = DataTag<int64_t>;
    using FloatTag = DataTag<float>;
    using DoubleTag = DataTag<double>;

    // Represents a generic array, used for StringTag,
    // ByteArrayTag, IntArrayTag, and LongArrayTag
    template <typename T>
    class ArrayTag : public DataTag<T*>
    {
    public:
        ArrayTag<T>(nbt_node* node) : DataTag<T*>(node) {}

        const T* GetValue() const;

        inline const T operator[](int32_t index) const { return GetValue()[index]; }

        // Methods for range-based for loops
        // TODO: check if this works for IntArrayTag and LongArrayTag
        const T* begin() { return GetValue(); }
        const T* end() { return GetValue() + Tag::GetLength(); }
    };

    // A string is an array of char. Its length is
    // not stored by cNBT, so Tag::GetLength uses strlen
    using StringTag = ArrayTag<char>;

    // Standard array types
    using ByteArrayTag = ArrayTag<int8_t>;
    using IntArrayTag = ArrayTag<int32_t>;
    using LongArrayTag = ArrayTag<int64_t>;

    // Represents a list of nameless tags of the same type
    template <typename T>
    class ListTag : public Tag
    {
        // Iterator for use in range-based for loops.
        // TODO: this might not need to be a nested class
        class iterator
        {
            list_head* m_pos;
        public:
            using iterator_category = std::input_iterator_tag;
            using value_type = T;
            using difference_type = int32_t;
            using pointer = const T*;
            using reference = const T&;

            // TODO: could inline some of this

            explicit iterator(ListTag<T>& tag, int32_t offset = 0) : m_pos(Internal::GetFirstItem(tag.m_node))
            {
                // If we have an offset, call operator++ that number of times
                for (int i = 0; i < offset; i++) { ++(*this); }
            }
            iterator& operator++() { m_pos = Internal::GetNextItem(m_pos); return *this; }
            bool operator==(iterator other) const { return m_pos == other.m_pos; }
            bool operator!=(iterator other) const { return !(*this == other); }
            T operator*() const; // defined below
        };

    public:
        ListTag<T>(nbt_node *node) : Tag(node) {}

        T Get(int32_t index) const;

        inline const T operator[](int32_t index) const { return Get(index); }

        iterator begin() { return iterator(*this); }
        iterator end() { return iterator(*this, Tag::GetLength()); }
    };

    // A Compound Tag is an unordered list of named tags
    class CompoundTag : public Tag
    {
        // TODO: see ListTag::iterator
        class iterator
        {
            list_head* m_pos;
        public:
            using iterator_category = std::input_iterator_tag;
            using value_type = Tag;
            using difference_type = int32_t;
            using pointer = const Tag*;
            using reference = const Tag&;

            // TODO: could inline some of this

            explicit iterator(CompoundTag& tag, int32_t offset = 0) : m_pos(Internal::GetFirstItem(tag.m_node))
            {
                // If we have an offset, call operator++ that number of times
                for (int i = 0; i < offset; i++) { ++(*this); }
            }
            iterator& operator++() { m_pos = Internal::GetNextItem(m_pos); return *this; }
            bool operator==(iterator other) const { return m_pos == other.m_pos; }
            bool operator!=(iterator other) const { return !(*this == other); }
            Tag operator*() const { return Tag::CreateTag(Internal::GetListEntry(m_pos)); }
        };


        // only true if the tree was allocated by constructor
        const bool m_isRoot;
        
        CompoundTag(nbt_node* node, bool isRoot) : Tag(node), m_isRoot(false) {}
    public:
        CompoundTag(nbt_node* node) : CompoundTag(node, false) {}

        // Read tree from file 
        CompoundTag(const char* filename);

        // Read tree from compressed data
        CompoundTag(const void* data, size_t length);

        ~CompoundTag();

        template <typename T>
        const ListTag<T> GetList(const char* name) const;

        template <typename T>
        const T Get(const char* name) const;

        inline const CompoundTag GetCompound(const char* name) const { return Get<CompoundTag>(name); }

        const Tag operator[](const char* name) const;

        iterator begin() { return iterator(*this); }
        iterator end() { return iterator(*this, Tag::GetLength()); }
    };

    //==================================================

    template <typename T>
    T ListTag<T>::iterator::operator*() const
    {
        if constexpr (std::is_base_of<Tag, T>::value) {
            return T(Internal::GetListEntry(m_pos));
        } else {
            // DataTag<T> can be auto converted to T
            return DataTag<T>(Internal::GetListEntry(m_pos));
        }
    }

    template <typename T>
    T ListTag<T>::Get(int32_t index) const
    {
        nbt_node* result = Internal::ListItem(m_node, index);
        if constexpr (std::is_base_of<Tag, T>::value) {
            return T(result);
        } else {
            // DataTag<T> can be auto converted to T
            return DataTag<T>(result);
        }
    }

    //==================================================

    template <typename T>
    const ListTag<T> CompoundTag::GetList(const char *name) const
    {
        nbt_node *result = Internal::FindByName(m_node, name);

        // TODO: do these checks in ListTag ctor
        if (result == NULL)
            return NULL;
        //if (result->type != TAG_LIST)
        //    return NULL;

        return ListTag<T>(result);
    }

    template <typename T>
    const T CompoundTag::Get(const char* name) const
    {
        nbt_node* result = Internal::FindByName(m_node, name);
        if constexpr (std::is_base_of<Tag, T>::value) {
            return T(result);
        }
        else {
            // DataTag<T> can be auto converted to T
            return DataTag<T>(result);
        }
    }
}