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

namespace Arbutils::Collections {
    template <class KeyT, class ValueT> class Dictionary {
    private:
        std::unordered_map<KeyT, ValueT> _map;

        using iterator = typename std::unordered_map<KeyT, ValueT, std::hash<KeyT>, std::equal_to<KeyT>,
                                                     std::allocator<std::pair<const KeyT, ValueT>>>::iterator;
    public:
        explicit Dictionary(size_t capacity) : _map(capacity) {}
        explicit Dictionary(const std::initializer_list<std::pair<const KeyT, ValueT>>& l) : _map(l){}

        inline void Insert(KeyT key, ValueT value) {
            auto v = _map.insert({key, value});
#ifndef NO_ASSERT
            if (!v.second)
                throw std::logic_error("Key already exists");
#endif
        }

        inline void Set(KeyT key, ValueT value) { _map[key] = value; }

        [[nodiscard]] inline ValueT& Get(KeyT key) {
#ifndef NO_ASSERT
            auto find = _map.find(key);
            if (find == _map.end()) {
                throw std::logic_error("Key not found");
            }
            return find->second;
#else
            return _map[key];
#endif
        }

        [[nodiscard]] inline const ValueT& Get(KeyT key) const {
#ifndef NO_ASSERT
            auto find = _map.find(key);
            if (find == _map.end()) {
                throw std::logic_error("Key not found");
            }
            return find->second;
#else
            return _map[key];
#endif
        }

        inline bool TryGet(KeyT key, ValueT& out) const {
            auto find = _map.find(key);
            if (find == _map.end()) {
                return false;
            }
            out = find->second;
            return true;
        }

        [[nodiscard]] inline size_t Count() const { return _map.size(); }

        inline bool Has(KeyT key) const noexcept { return _map.find(key) != _map.end(); }

        inline ValueT& operator[](KeyT key) { return Get(key); }
        inline const ValueT& operator[](KeyT key) const { return Get(key); }

        iterator begin() { return _map.begin(); }
        iterator begin() const { return _map.begin(); }

        iterator end() { return _map.end(); }
        iterator end() const { return _map.end(); }
    };
}

#endif // ARBUTILS_DICTIONARY_HPP