#ifndef HASHMAP_H
#define HASHMAP_H

#include "utils/Array.h"
#include "utils/List.h"
#include "utils/StringBuffer.h"
#include "utils/Types.h"
#include "utils/Utils.h"

#include <unordered_map>

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

    public:
        V value;

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

    private:
        int next;

        Node(const K& key, const V& value) : key(key), value(value), next(-1) {
        }

        Node(const K& key, V&& value)
            : key(key), value(std::move(value)), next(-1) {
        }

        template<typename... Args>
        Node(const K& key, Args&&... args)
            : key(key), value(std::forward<Args>(args)...), next(-1) {
        }
    };

    template<typename N, typename R>
    class BaseIterator {
        N& nodes;
        int indexA;
        int indexB;

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

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

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

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

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

    typedef BaseIterator<List<List<Node>>, Node> Iterator;
    typedef BaseIterator<const List<List<Node>>, const Node> ConstIterator;

private:
    List<List<Node>> nodes;
    int elements;

public:
    HashMap(int minCapacity = 8) : elements(0) {
        nodes.resize(1 << Utils::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;
    }

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

    HashMap& add(const K& key, V&& value) {
        rehash();
        Hash h = hash(key);
        V* v = searchList(key, h);
        if(v == nullptr) {
            nodes[h].add(key, std::move(value));
            elements++;
        } else {
            *v = std::move(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.remove(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;
    }

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

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

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

    ConstIterator end() const {
        return ConstIterator(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;
    }

    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