export module Core.HashMap;

export import Core.Types;
export import Core.Utility;
export import Core.New;

import Core.List;
import Core.ToString;
import Core.AlignedData;
import Core.Math;
import Core.Meta;

export template<typename T>
concept Hashable = requires(const T& t) { t.hashCode(); };
export template<typename T>
concept HashCast = requires(const T& t) { static_cast<size_t>(t); };

export namespace Core {
    template<Hashable H>
    inline size_t hashCode(const H& key) {
        return key.hashCode();
    }

    template<HashCast H>
    inline size_t hashCode(const H& key) {
        return static_cast<size_t>(key);
    }

    template<Moveable K, Moveable V>
    struct HashMap final {
        template<typename Value>
        class Node final {
            friend HashMap;
            K key;

        public:
            Value& value;

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

            size_t toString(StringBase& b) const noexcept {
                return b.addFormat("{} = {}", key, value);
            }

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

    private:
        static inline K INVALID = {};

        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,
                bool invalidSet) noexcept :
                currentKey(key), endKey(endKey_), currentValue(value) {
                if(!invalidSet) {
                    skip();
                }
            }

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

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

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

        private:
            void skip() noexcept {
                while(currentKey != endKey && *currentKey == INVALID) {
                    ++currentKey;
                    ++currentValue;
                }
            }
        };

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

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

        static const K& accessKey(const K& key, const V&) noexcept {
            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 noexcept {
                return {
                    map.keys.begin(), map.keys.end(), map.values,
                    map.invalidSet};
            }

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

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

        using ConstKeyIteratorAdapter =
            IteratorAdapter<const HashMap, ConstKeyIterator>;

    private:
        List<K> keys{};
        V* values = nullptr;
        List<i8> jumps{};
        size_t entries = 0;
        bool invalidSet = false;

    public:
        HashMap() noexcept = default;

        HashMap(const HashMap& other) noexcept {
            for(const auto& e : other) {
                add(e.getKey(), e.value);
            }
        }

        HashMap(HashMap&& other) noexcept {
            swap(other);
        }

        ~HashMap() noexcept {
            size_t length = keys.getLength();
            if(length > 0) {
                for(size_t i = 1; i < length; i++) {
                    if(keys[i] != INVALID) {
                        values[i].~V();
                    }
                }
                if(invalidSet) {
                    values[length].~V();
                }
            }
            deleteWithSourceN<AlignedType<V>>(
                reinterpret_cast<AlignedType<V>*>(values));
        }

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

        void rehash(size_t minCapacity) noexcept {
            if(minCapacity <= keys.getLength()) {
                return;
            }
            HashMap<K, V> map;
            size_t l = (1lu << roundUpLog2(max(minCapacity, 8lu))) + 1;
            map.keys.resize(l, INVALID);
            map.values =
                reinterpret_cast<V*>(newWithSourceN<AlignedType<V>>(l));
            map.jumps.resize(l, 0);
            size_t length = keys.getLength();
            if(length > 0) {
                for(size_t i = 1; i < length; i++) {
                    if(keys[i] != INVALID) {
                        map.add(keys[i], Core::move(values[i]));
                    }
                }
                if(invalidSet) {
                    map.add(INVALID, Core::move(values[length]));
                }
            }
            swap(map);
        }

        template<typename... Args>
        bool tryEmplace(V*& v, const K& key, Args&&... args) noexcept {
            size_t index = 0;
            if(key == INVALID) {
                if(invalidSet) {
                    return false;
                }
                rehash(1);
                invalidSet = true;
            } else {
                index = searchSlot(key);
                if(keys[index] == key) {
                    return false;
                }
            }
            keys[index] = key;
            static_assert(
                noexcept(new(values + index) V(Core::forward<Args>(args)...)));
            v = new(values + index) V(Core::forward<Args>(args)...);
            entries++;
            markSlot(key);
            return true;
        }

        template<typename VA>
        V& put(const K& key, VA&& value) noexcept {
            size_t index = 0;
            if(key == INVALID) {
                if(invalidSet) {
                    static_assert(
                        noexcept(values[0] = Core::forward<VA>(value)));
                    return (values[0] = Core::forward<VA>(value));
                }
                rehash(1);
                invalidSet = true;
            } else {
                index = searchSlot(key);
                if(keys[index] == key) {
                    static_assert(
                        noexcept(values[index] = Core::forward<VA>(value)));
                    return (values[index] = Core::forward<VA>(value));
                }
            }
            static_assert(
                noexcept(new(values + index) V(Core::forward<VA>(value))));
            new(values + index) V(Core::forward<VA>(value));
            entries++;
            keys[index] = key;
            markSlot(key);
            return values[index];
        }

        template<typename VA>
        HashMap& add(const K& key, VA&& value) noexcept {
            put(key, Core::forward<VA>(value));
            return *this;
        }

        bool remove(const K& key) noexcept {
            size_t index = 0;
            if(key == INVALID) {
                if(!invalidSet) {
                    return false;
                }
                invalidSet = false;
            } else {
                index = searchSlot(key);
                if(keys[index] != key) {
                    return false;
                }
            }
            values[index].~V();
            entries--;
            demarkSlot(key);
            keys[index] = INVALID;
            return true;
        }

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

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

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

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

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

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

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

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

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

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

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

        void swap(HashMap& o) noexcept {
            Core::swap(o.keys, keys);
            Core::swap(o.values, values);
            Core::swap(o.jumps, jumps);
            Core::swap(o.entries, entries);
            Core::swap(o.invalidSet, invalidSet);
        }

    private:
        // clang-format off
        #define FOR_EACH_HASH_START()                           \
            do {                                                \
                size_t baseHash = hashCode(key) * 514'685'581u; \
                size_t end = keys.getLength() - 2;              \
                for(size_t i = 0; i <= 5; i++) {                \
                    size_t hash = 1 + ((baseHash + i) & end)
        #define FOR_EACH_HASH_STOP() \
                }                    \
            } while(false)
        // clang-format on

        size_t searchSlot(const K& key) noexcept {
            rehash(1);
            while(true) {
                // rehash on bad clustering
                FOR_EACH_HASH_START();
                if((keys[hash] == INVALID && jumps[hash] == 0) ||
                   keys[hash] == key) {
                    return hash;
                }
                FOR_EACH_HASH_STOP();
                rehash(keys.getLength() + 1);
            }
        }

        void markSlot(const K& key) noexcept {
            FOR_EACH_HASH_START();
            if(keys[hash] == key) {
                return;
            }
            jumps[hash]++;
            FOR_EACH_HASH_STOP();
        }

        void demarkSlot(const K& key) noexcept {
            FOR_EACH_HASH_START();
            if(keys[hash] == key) {
                return;
            }
            jumps[hash]--;
            FOR_EACH_HASH_STOP();
        }

        template<typename Value>
        Value* searchValue(const K& key) const noexcept {
            if(keys.getLength() != 0) {
                if(key == INVALID) {
                    return invalidSet ? values : nullptr;
                }
                FOR_EACH_HASH_START();
                if(keys[hash] == key) {
                    return values + hash;
                } else if(jumps[hash] == 0) {
                    return nullptr;
                }
                FOR_EACH_HASH_STOP();
            }
            return nullptr;
        }
    };
}