#ifndef CORE_PROBING_HPPASHMAP_HPP
#define CORE_PROBING_HPPASHMAP_HPP

#include "core/data/List.hpp"
#include "core/utils/AlignedData.hpp"
#include "core/utils/ArrayString.hpp"
#include "core/utils/HashCode.hpp"
#include "core/utils/Logger.hpp"

namespace Core {
    template<typename K, typename V>
    struct ProbingHashMap final {
        template<typename Value>
        class Node final {
            friend ProbingHashMap;
            friend List<Node>;
            K key;

        public:
            Value& value;

            const K& getKey() const {
                return key;
            }

            template<typename String>
            check_return Error toString(String& s) const {
                CORE_RETURN_ERROR(s.append(key));
                CORE_RETURN_ERROR(s.append(" = "));
                return s.append(value);
            }

        private:
            Node(const K& key_, Value& value_) : key(key_), value(value_) {
            }
        };

    private:
        template<typename Value, typename R, R (*A)(const K&, Value&)>
        class Iterator final {
            const K* currentKey;
            const K* endKey;
            Value* currentValue;

        public:
            Iterator(const K* key, const K* endKey_, Value* value)
                : currentKey(key), endKey(endKey_), currentValue(value) {
                skip();
            }

            Iterator& operator++() {
                ++currentKey;
                ++currentValue;
                skip();
                return *this;
            }

            bool operator!=(const Iterator& other) const {
                return currentKey != other.currentKey;
            }

            R operator*() const {
                return A(*currentKey, *currentValue);
            }

        private:
            void skip() {
                while(currentKey != endKey && *currentKey == emptyValue<K>()) {
                    ++currentKey;
                    ++currentValue;
                }
            }
        };

        template<typename Value>
        static Node<Value> access(const K& key, Value& value) {
            return Node<Value>(key, value);
        }

        template<typename Value>
        static Value& accessValue(const K&, Value& value) {
            return value;
        }

        static const K& accessKey(const K& key, const V&) {
            return key;
        }

        template<typename Value>
        using BaseEntryIterator = Iterator<Value, Node<Value>, access<Value>>;
        using EntryIterator = BaseEntryIterator<V>;
        using ConstEntryIterator = BaseEntryIterator<const V>;

        template<typename Value>
        using BaseValueIterator = Iterator<Value, Value&, accessValue<Value>>;
        using ValueIterator = BaseValueIterator<V>;
        using ConstValueIterator = BaseValueIterator<const V>;

        using ConstKeyIterator = Iterator<const V, const K&, accessKey>;

        template<typename M, typename I>
        struct IteratorAdapter final {
            M& map;

            I begin() const {
                return {map.keys.begin(), map.keys.end(), map.values};
            }

            I end() const {
                return {map.keys.end(), map.keys.end(), nullptr};
            }
        };

        using ValueIteratorAdapter =
            IteratorAdapter<ProbingHashMap, ValueIterator>;
        using ConstValueIteratorAdapter =
            IteratorAdapter<const ProbingHashMap, ConstValueIterator>;

        using ConstKeyIteratorAdapter =
            IteratorAdapter<const ProbingHashMap, ConstKeyIterator>;

    private:
        List<K> keys;
        V* values;
        i64 entries;

    public:
        ProbingHashMap() : keys(), values(nullptr), entries(0) {
        }

        ProbingHashMap(const ProbingHashMap& other) = delete;

        ProbingHashMap(ProbingHashMap&& other) : ProbingHashMap() {
            swap(other);
        }

        ~ProbingHashMap() {
            for(i64 i = 0; i < keys.getLength(); i++) {
                if(keys[i] != emptyValue<K>()) {
                    values[i].~V();
                }
            }
            delete[] reinterpret_cast<AlignedType<V>*>(values);
        }

        ProbingHashMap& operator=(const ProbingHashMap& other) = delete;

        ProbingHashMap& operator=(ProbingHashMap&& other) {
            swap(other);
            return *this;
        }

        check_return Error copyFrom(const ProbingHashMap& other) {
            ProbingHashMap copy;
            for(const auto& e : other) {
                CORE_RETURN_ERROR(copy.add(e.getKey(), e.value));
            }
            swap(copy);
            return Error::NONE;
        }

        check_return Error rehash(i64 minCapacity) {
            if(minCapacity <= keys.getLength()) {
                return Error::NONE;
            }
            ProbingHashMap<K, V> map;
            i64 l = 1l << Math::roundUpLog2(Core::Math::max(minCapacity, 8l));
            CORE_RETURN_ERROR(map.keys.resize(l, emptyValue<K>()));
            map.values = reinterpret_cast<V*>(new AlignedType<V>[l]);
            if(map.values == nullptr) {
                return Error::OUT_OF_MEMORY;
            }
            for(i64 i = 0; i < keys.getLength(); i++) {
                if(keys[i] != emptyValue<K>()) {
                    CORE_RETURN_ERROR(map.add(keys[i], values[i]));
                }
            }
            swap(map);
            return Error::NONE;
        }

        template<typename... Args>
        check_return Error tryEmplace(V*& v, const K& key, Args&&... args) {
            if(key == emptyValue<K>()) {
                return Error::INVALID_ARGUMENT;
            }
            i64 index = 0;
            CORE_RETURN_ERROR(searchSlot(index, key));
            if(keys[index] == key) {
                return Error::EXISTING_KEY;
            }
            keys[index] = key;
            v = new(values + index) V(Core::forward<Args>(args)...);
            entries++;
            return Error::NONE;
        }

        template<typename VA>
        check_return Error put(V*& v, const K& key, VA&& value) {
            if(key == emptyValue<K>()) {
                return Error::INVALID_ARGUMENT;
            }
            i64 index = 0;
            CORE_RETURN_ERROR(searchSlot(index, key));
            if(keys[index] == key) {
                values[index] = Core::forward<VA>(value);
            } else {
                new(values + index) V(Core::forward<VA>(value));
                entries++;
            }
            keys[index] = key;
            v = reinterpret_cast<V*>(values + index);
            return Error::NONE;
        }

        template<typename VA>
        check_return Error add(const K& key, VA&& value) {
            V* v = nullptr;
            return put(v, key, Core::forward<VA>(value));
        }

        const V* search(const K& key) const {
            return searchValue<const V>(key);
        }

        V* search(const K& key) {
            return searchValue<V>(key);
        }

        bool contains(const K& key) const {
            return search(key) != nullptr;
        }

        ProbingHashMap& clear() {
            ProbingHashMap<K, V> map;
            swap(map);
            return *this;
        }

        ConstKeyIteratorAdapter getKeys() const {
            return {*this};
        }

        ValueIteratorAdapter getValues() {
            return {*this};
        }

        ConstValueIteratorAdapter getValues() const {
            return {*this};
        }

        EntryIterator begin() {
            return {keys.begin(), keys.end(), values};
        }

        EntryIterator end() {
            return {keys.end(), keys.end(), nullptr};
        }

        ConstEntryIterator begin() const {
            return {keys.begin(), keys.end(), values};
        }

        ConstEntryIterator end() const {
            return {keys.end(), keys.end(), nullptr};
        }

        template<typename String>
        check_return Error toString(String& s) const {
            return Core::toString(s, *this);
        }

        void swap(ProbingHashMap& o) {
            Core::swap(o.keys, keys);
            Core::swap(o.values, values);
            Core::swap(o.entries, entries);
        }

    private:
        check_return Error searchSlot(i64& slot, const K& key) {
            i64 rehashFactor = 2;
            while(true) {
                CORE_RETURN_ERROR(rehash(entries * rehashFactor + 1));
                u64 baseHash = hashCode(key) * 514685581u;
                u64 end = static_cast<u64>(keys.getLength()) - 1;
                // rehash on bad clustering
                for(u64 i = 0; i <= 5; i++) {
                    i64 hash = static_cast<i64>((baseHash + i) & end);
                    if(keys[hash] == emptyValue<K>() || keys[hash] == key) {
                        slot = hash;
                        return Core::Error::NONE;
                    }
                }
                rehashFactor *= 2;
            }
        }

        template<typename Value>
        Value* searchValue(const K& key) const {
            if(keys.getLength() != 0) {
                u64 baseHash = hashCode(key) * 514685581u;
                u64 end = static_cast<u64>(keys.getLength()) - 1;
                for(u32 i = 0; i <= end; i++) [[unlikely]] {
                    i64 hash = static_cast<i64>((baseHash + i) & end);
                    if(keys[hash] == key) [[likely]] {
                        return values + hash;
                    } else if(keys[hash] == emptyValue<K>()) {
                        return nullptr;
                    }
                }
            }
            return nullptr;
        }
    };
}

#endif