#ifndef HASHMAP_H
#define HASHMAP_H

#include "data/List.h"
#include "math/Math.h"
#include "utils/StringBuffer.h"
#include "utils/Types.h"

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

    public:
        V value;

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

    private:
        template<typename... Args>
        Node(const K& key, Args&&... args)
            : key(key), value(std::forward<Args>(args)...) {
        }
    };
    using Nodes = List<List<Node>>;

    template<typename N, typename I, typename R, R& (*A)(I&)>
    class Iterator final {
        N& nodes;
        int indexA;
        int indexB;

    public:
        Iterator(N& nodes, int indexA, int indexB)
            : nodes(nodes), indexA(indexA), indexB(indexB) {
            skip();
        }

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

        bool operator!=(const Iterator& other) const {
            return indexA != other.indexA || indexB != other.indexB;
        }

        R& operator*() const {
            return A(nodes[indexA][indexB]);
        }

    private:
        void skip() {
            while(indexA < nodes.getLength() &&
                  indexB >= nodes[indexA].getLength()) {
                indexA++;
                indexB = 0;
            }
        }
    };

    template<typename R>
    static R& access(R& node) {
        return node;
    }

    template<typename I, typename R>
    static R& accessValue(I& node) {
        return node.value;
    }

    static const K& accessKey(const Node& node) {
        return node.getKey();
    }

    template<typename N, typename R>
    using BaseEntryIterator = Iterator<N, R, R, access<R>>;
    using EntryIterator = BaseEntryIterator<Nodes, Node>;
    using ConstEntryIterator = BaseEntryIterator<const Nodes, const Node>;

    template<typename N, typename I, typename R>
    using BaseValueIterator = Iterator<N, I, R, accessValue<I, R>>;
    using ValueIterator = BaseValueIterator<Nodes, Node, V>;
    using ConstValueIterator =
        BaseValueIterator<const Nodes, const Node, const V>;

    using ConstKeyIterator =
        Iterator<const Nodes, const Node, const K, accessKey>;

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

        I begin() const {
            return I(map.nodes, 0, 0);
        }

        I end() const {
            return I(map.nodes, map.nodes.getLength(), 0);
        }
    };

    using EntryIteratorAdapter = IteratorAdapter<HashMap, EntryIterator>;
    using ConstEntryIteratorAdapter =
        IteratorAdapter<const HashMap, ConstEntryIterator>;

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

    using ConstKeyIteratorAdapter =
        IteratorAdapter<const HashMap, ConstKeyIterator>;

private:
    Nodes nodes;
    int elements;

public:
    HashMap(int minCapacity = 8) : elements(0) {
        nodes.resize(1 << Math::roundUpLog2(minCapacity));
    }

    template<typename... Args>
    bool tryEmplace(const K& key, Args&&... args) {
        rehash();
        Hash h = hash(key);
        V* v = searchList(key, h);
        if(v == nullptr) {
            nodes[h].add(key, std::forward<Args>(args)...);
            elements++;
            return false;
        }
        return true;
    }

    template<typename VA>
    HashMap& add(const K& key, VA&& value) {
        rehash();
        Hash h = hash(key);
        V* v = searchList(key, h);
        if(v == nullptr) {
            nodes[h].add(key, std::forward<VA>(value));
            elements++;
        } else {
            *v = std::forward<VA>(value);
        }
        return *this;
    }

    bool remove(const K& key) {
        List<Node>& list = nodes[hash(key)];
        for(int i = 0; i < list.getLength(); i++) {
            if(list[i].key == key) {
                list.removeBySwap(i);
                return true;
            }
        }
        return false;
    }

    const V* search(const K& key) const {
        return searchList(key, hash(key));
    }

    V* search(const K& key) {
        return searchList(key, hash(key));
    }

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

    HashMap& clear() {
        for(List<Node>& n : nodes) {
            n.clear();
        }
        elements = 0;
        return *this;
    }

    EntryIteratorAdapter entries() {
        return {*this};
    }

    ConstEntryIteratorAdapter entries() const {
        return {*this};
    }

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

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

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

    EntryIterator begin() {
        return EntryIterator(nodes, 0, 0);
    }

    EntryIterator end() {
        return EntryIterator(nodes, nodes.getLength(), 0);
    }

    ConstEntryIterator begin() const {
        return ConstEntryIterator(nodes, 0, 0);
    }

    ConstEntryIterator end() const {
        return ConstEntryIterator(nodes, nodes.getLength(), 0);
    }

    template<int L>
    void toString(StringBuffer<L>& s) const {
        s.append("[");
        bool c = false;
        for(const List<Node>& list : nodes) {
            for(const Node& n : list) {
                if(c) {
                    s.append(", ");
                }
                s.append(n.key).append(" = ").append(n.value);
                c = true;
            }
        }
        s.append("]");
    }

private:
    template<typename H>
    Hash hash(const H& key) const {
        return fullHash(key) & (nodes.getLength() - 1);
    }

    template<typename H>
    Hash fullHash(const H& key) const {
        return key.hashCode();
    }

    Hash fullHash(int key) const {
        return key;
    }

    Hash fullHash(unsigned int key) const {
        return key;
    }

    void rehash() {
        if(elements < nodes.getLength()) {
            return;
        }
        HashMap<K, V> map(nodes.getLength() * 2);
        for(List<Node>& list : nodes) {
            for(Node& n : list) {
                map.tryEmplace(n.key, std::move(n.value));
            }
        }
        *this = std::move(map);
    }

    const V* searchList(const K& key, Hash h) const {
        for(const Node& n : nodes[h]) {
            if(n.key == key) {
                return &n.value;
            }
        }
        return nullptr;
    }

    V* searchList(const K& key, Hash h) {
        return const_cast<V*>(
            static_cast<const HashMap*>(this)->searchList(key, h));
    }
};

#endif