#ifndef ARBUTILS_DICTIONARY_HPP
#define ARBUTILS_DICTIONARY_HPP
#include <cstdint>
#include <unordered_map>
#include "../Assert.hpp"

namespace ArbUt {
    /// @brief Wrapper around unordered_map, allowing safer access and adding several helper methods.
    template <class KeyT, class ValueT> class Dictionary {
    private:
        std::unordered_map<KeyT, ValueT> _map;

        using iterator = typename std::unordered_map<KeyT, ValueT>::iterator;
        using const_iterator = typename std::unordered_map<KeyT, ValueT>::const_iterator;

    public:
        Dictionary() : _map() {}
        /// @brief Initialises a dictionary with a certain capacity.
        explicit Dictionary(size_t capacity) : _map(capacity) {}
        /// @brief Initialises a dictionary from an initializer_list.
        Dictionary(const std::initializer_list<std::pair<const KeyT, ValueT>>& l) : _map(l) {}

        /// @brief Removes all items from the dictionary.
        inline void Clear() noexcept { _map.clear(); }

        /// @brief Inserts a new item in the dictionary. This will throw if the dictionary already contains the key.
        inline void Insert(const KeyT& key, const ValueT& value) {
            [[maybe_unused]] const auto& v = _map.insert({key, value});
#ifndef NO_ASSERT
            if (!v.second)
                throw ArbUt::Exception("Key already exists");
#endif
        }

        /// @brief Sets a key in the dictionary to a specific value.
        inline void Set(const KeyT& key, const ValueT& value) { _map[key] = value; }

        /// @brief Gets a value from the dictionary.
        [[nodiscard]] inline ValueT& Get(const KeyT& key) {
#ifndef NO_ASSERT
            return _map.at(key);
#else
            return _map[key];
#endif
        }
        /// @brief Gets a value from the dictionary.
        [[nodiscard]] inline const ValueT& Get(const KeyT& key) const { return _map.at(key); }

        /// @brief Try to get an item from the dictionary using a key. Returns false if no item is found, and out will
        /// not be touched in that case.
        inline bool TryGet(const KeyT& key, ValueT& out) const noexcept {
            const auto& find = _map.find(key);
            if (find == _map.end()) {
                return false;
            }
            out = find->second;
            return true;
        }

        /// @brief Removes an item with a certain key from the dictionary
        inline void Remove(const KeyT& key) { _map.erase(key); }

        /// @brief Returns the number of items in the dictionary.
        [[nodiscard]] inline size_t Count() const noexcept { return _map.size(); }

        /// @brief Checks whether the dictionary contains a specific key.
        inline bool Has(const KeyT& key) const noexcept { return _map.find(key) != _map.end(); }

        /// @brief Indexing operator to get a value from the dictionary using a key.
        inline ValueT& operator[](const KeyT& key) { return Get(key); }
        /// @brief Indexing operator to get a value from the dictionary using a key.
        inline const ValueT& operator[](const KeyT& key) const { return Get(key); }

        /// @brief returns an iterator to the beginning of the specified bucket
        iterator begin() noexcept { return _map.begin(); }
        /// @brief returns an iterator to the beginning of the specified bucket
        const_iterator begin() const noexcept { return _map.begin(); }

        /// @brief returns an iterator to the end of the specified bucket
        iterator end() noexcept { return _map.end(); }
        /// @brief returns an iterator to the end of the specified bucket
        const_iterator end() const { return _map.end(); }

        /// @brief returns the backing unordered_map.
        const std::unordered_map<KeyT, ValueT>& GetStdMap() const noexcept { return _map; }
        /// @brief returns the backing unordered_map.
        std::unordered_map<KeyT, ValueT>& GetStdMap() noexcept { return _map; }
    };
}

#endif // ARBUTILS_DICTIONARY_HPP