#ifndef ARBUTILS_STRINGVIEWDICTIONARY_HPP
#define ARBUTILS_STRINGVIEWDICTIONARY_HPP
#include <algorithm>
#include <functional>
#include <optional>
#include <unordered_map>
#include "../Ensure.hpp"
#include "../StringView.hpp"

namespace ArbUt {
    /// @brief Wrapper around unordered_map, allowing safer access and adding several helper methods.
    template <class ValueT> class StringViewDictionary {
    public:
        /// @brief Custom hash class to allow for accessing values by both StringView and direct hash
        struct StringViewHash {
            /// @brief The actual hash
            using hash_type = std::hash<StringView>;
            /// @brief As stolen from the documentation
            using is_transparent = void;

            /// @brief Support hashing through a StringView
            size_t operator()(StringView const& str) const { return hash_type{}(str); }
            /// @brief Support hashing through the direct hash access.
            size_t operator()(u32 hash) const { return hash; }
        };

        /// @brief The underlying std map used
        using map_type = typename std::unordered_map<StringView, ValueT, StringViewHash, std::equal_to<>>;
        /// @brief The iterator type of the Dictionary
        using iterator = typename map_type::iterator;
        /// @brief The const iterator type of the Dictionary
        using const_iterator = typename map_type::const_iterator;

    private:
        map_type _map;

    public:
        /// @brief The type used for the indexer of the dictionary.
        typedef StringView keyType;
        /// @brief The type used for the value in the dictionary.
        typedef ValueT valueType;

        StringViewDictionary() : _map() {}
        /// @brief Initialises a dictionary with a certain capacity.
        explicit StringViewDictionary(size_t capacity) : _map(capacity) {}
        /// @brief Initialises a dictionary from an initializer_list.
        StringViewDictionary(const std::initializer_list<std::pair<const StringView, ValueT>>& l) : _map(l) {}
        /// @brief Copy operator
        StringViewDictionary(const StringViewDictionary& other) noexcept : _map(other._map) {}
        /// @brief Assignment operator
        StringViewDictionary& operator=(const StringViewDictionary& other) noexcept {
            if (this == &other) {
                return *this;
            }
            _map = other._map;
            return *this;
        }

        /// @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 StringView& 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 StringView& key, const ValueT& value) { _map[key] = value; }

        /// @brief Gets a value from the dictionary.
        [[nodiscard]] inline ValueT& Get(const StringView& 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 StringView& key) const { return _map.at(key); }

        /// @brief Gets a value from the dictionary using the direct hash of the key.
        [[nodiscard]] inline ValueT& GetFromHash(u32 hash) {
#ifndef NO_ASSERT
            return _map.find(hash)->second;
#else
            return _map[key];
#endif
        }
        /// @brief Gets a value from the dictionary using the direct hash of the key.
        [[nodiscard]] inline const ValueT& GetFromHash(u32 hash) const { return _map.find(hash)->second; }

        /// @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 std::optional<std::reference_wrapper<const ValueT>> TryGet(const StringView& key) const noexcept {
            const auto& find = _map.find(key);
            if (find == _map.end()) {
                return {};
            }
            return std::ref(find->second);
        }

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

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

        /// @brief Removes an item with a certain key hash from the dictionary
        inline void Remove(u32 keyHash) {
            auto find = _map.find(keyHash);
            if (find != _map.end()) {
                _map.erase(find->first);
            }
        }

        /// @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 StringView& 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 StringView& key) { return Get(key); }
        /// @brief Indexing operator to get a value from the dictionary using a key.
        inline const ValueT& operator[](const StringView& 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 map_type& GetStdMap() const noexcept { return _map; }
        /// @brief returns the backing unordered_map.
        map_type& GetStdMap() noexcept { return _map; }
    };
}

#endif // ARBUTILS_STRINGVIEWDICTIONARY_HPP