#ifndef XRPL_BASICS_PARTITIONED_UNORDERED_MAP_H
#define XRPL_BASICS_PARTITIONED_UNORDERED_MAP_H

#include <xrpl/beast/hash/uhash.h>
#include <xrpl/beast/utility/instrumentation.h>

#include <functional>
#include <optional>
#include <string>
#include <thread>
#include <unordered_map>
#include <utility>
#include <vector>

namespace ripple {

template <typename Key>
static std::size_t
extract(Key const& key)
{
    return key;
}

template <>
inline std::size_t
extract(std::string const& key)
{
    return ::beast::uhash<>{}(key);
}

template <
    typename Key,
    typename Value,
    typename Hash,
    typename Pred = std::equal_to<Key>,
    typename Alloc = std::allocator<std::pair<Key const, Value>>>
class partitioned_unordered_map
{
    std::size_t partitions_;

public:
    using key_type = Key;
    using mapped_type = Value;
    using value_type = std::pair<Key const, mapped_type>;
    using size_type = std::size_t;
    using difference_type = std::size_t;
    using hasher = Hash;
    using key_equal = Pred;
    using allocator_type = Alloc;
    using reference = value_type&;
    using const_reference = value_type const&;
    using pointer = value_type*;
    using const_pointer = value_type const*;
    using map_type = std::
        unordered_map<key_type, mapped_type, hasher, key_equal, allocator_type>;
    using partition_map_type = std::vector<map_type>;

    struct iterator
    {
        using iterator_category = std::forward_iterator_tag;
        partition_map_type* map_{nullptr};
        typename partition_map_type::iterator ait_;
        typename map_type::iterator mit_;

        iterator() = default;

        iterator(partition_map_type* map) : map_(map)
        {
        }

        reference
        operator*() const
        {
            return *mit_;
        }

        pointer
        operator->() const
        {
            return &(*mit_);
        }

        void
        inc()
        {
            ++mit_;
            while (mit_ == ait_->end())
            {
                ++ait_;
                if (ait_ == map_->end())
                    return;
                mit_ = ait_->begin();
            }
        }

        // ++it
        iterator&
        operator++()
        {
            inc();
            return *this;
        }

        // it++
        iterator
        operator++(int)
        {
            iterator tmp(*this);
            inc();
            return tmp;
        }

        friend bool
        operator==(iterator const& lhs, iterator const& rhs)
        {
            return lhs.map_ == rhs.map_ && lhs.ait_ == rhs.ait_ &&
                lhs.mit_ == rhs.mit_;
        }

        friend bool
        operator!=(iterator const& lhs, iterator const& rhs)
        {
            return !(lhs == rhs);
        }
    };

    struct const_iterator
    {
        using iterator_category = std::forward_iterator_tag;

        partition_map_type* map_{nullptr};
        typename partition_map_type::iterator ait_;
        typename map_type::iterator mit_;

        const_iterator() = default;

        const_iterator(partition_map_type* map) : map_(map)
        {
        }

        const_iterator(iterator const& orig)
        {
            map_ = orig.map_;
            ait_ = orig.ait_;
            mit_ = orig.mit_;
        }

        const_reference
        operator*() const
        {
            return *mit_;
        }

        const_pointer
        operator->() const
        {
            return &(*mit_);
        }

        void
        inc()
        {
            ++mit_;
            while (mit_ == ait_->end())
            {
                ++ait_;
                if (ait_ == map_->end())
                    return;
                mit_ = ait_->begin();
            }
        }

        // ++it
        const_iterator&
        operator++()
        {
            inc();
            return *this;
        }

        // it++
        const_iterator
        operator++(int)
        {
            const_iterator tmp(*this);
            inc();
            return tmp;
        }

        friend bool
        operator==(const_iterator const& lhs, const_iterator const& rhs)
        {
            return lhs.map_ == rhs.map_ && lhs.ait_ == rhs.ait_ &&
                lhs.mit_ == rhs.mit_;
        }

        friend bool
        operator!=(const_iterator const& lhs, const_iterator const& rhs)
        {
            return !(lhs == rhs);
        }
    };

private:
    std::size_t
    partitioner(Key const& key) const
    {
        return extract(key) % partitions_;
    }

    template <class T>
    static void
    end(T& it)
    {
        it.ait_ = it.map_->end();
        it.mit_ = it.map_->back().end();
    }

    template <class T>
    static void
    begin(T& it)
    {
        for (it.ait_ = it.map_->begin(); it.ait_ != it.map_->end(); ++it.ait_)
        {
            if (it.ait_->begin() == it.ait_->end())
                continue;
            it.mit_ = it.ait_->begin();
            return;
        }
        end(it);
    }

public:
    partitioned_unordered_map(
        std::optional<std::size_t> partitions = std::nullopt)
    {
        // Set partitions to the number of hardware threads if the parameter
        // is either empty or set to 0.
        partitions_ = partitions && *partitions
            ? *partitions
            : std::thread::hardware_concurrency();
        map_.resize(partitions_);
        XRPL_ASSERT(
            partitions_,
            "ripple::partitioned_unordered_map::partitioned_unordered_map : "
            "nonzero partitions");
    }

    std::size_t
    partitions() const
    {
        return partitions_;
    }

    partition_map_type&
    map()
    {
        return map_;
    }

    iterator
    begin()
    {
        iterator it(&map_);
        begin(it);
        return it;
    }

    const_iterator
    cbegin() const
    {
        const_iterator it(&map_);
        begin(it);
        return it;
    }

    const_iterator
    begin() const
    {
        return cbegin();
    }

    iterator
    end()
    {
        iterator it(&map_);
        end(it);
        return it;
    }

    const_iterator
    cend() const
    {
        const_iterator it(&map_);
        end(it);
        return it;
    }

    const_iterator
    end() const
    {
        return cend();
    }

private:
    template <class T>
    void
    find(key_type const& key, T& it) const
    {
        it.ait_ = it.map_->begin() + partitioner(key);
        it.mit_ = it.ait_->find(key);
        if (it.mit_ == it.ait_->end())
            end(it);
    }

public:
    iterator
    find(key_type const& key)
    {
        iterator it(&map_);
        find(key, it);
        return it;
    }

    const_iterator
    find(key_type const& key) const
    {
        const_iterator it(&map_);
        find(key, it);
        return it;
    }

    template <class T, class U>
    std::pair<iterator, bool>
    emplace(std::piecewise_construct_t const&, T&& keyTuple, U&& valueTuple)
    {
        auto const& key = std::get<0>(keyTuple);
        iterator it(&map_);
        it.ait_ = it.map_->begin() + partitioner(key);
        auto [eit, inserted] = it.ait_->emplace(
            std::piecewise_construct,
            std::forward<T>(keyTuple),
            std::forward<U>(valueTuple));
        it.mit_ = eit;
        return {it, inserted};
    }

    template <class T, class U>
    std::pair<iterator, bool>
    emplace(T&& key, U&& val)
    {
        iterator it(&map_);
        it.ait_ = it.map_->begin() + partitioner(key);
        auto [eit, inserted] =
            it.ait_->emplace(std::forward<T>(key), std::forward<U>(val));
        it.mit_ = eit;
        return {it, inserted};
    }

    void
    clear()
    {
        for (auto& p : map_)
            p.clear();
    }

    iterator
    erase(const_iterator position)
    {
        iterator it(&map_);
        it.ait_ = position.ait_;
        it.mit_ = position.ait_->erase(position.mit_);

        while (it.mit_ == it.ait_->end())
        {
            ++it.ait_;
            if (it.ait_ == it.map_->end())
                break;
            it.mit_ = it.ait_->begin();
        }

        return it;
    }

    std::size_t
    size() const
    {
        std::size_t ret = 0;
        for (auto& p : map_)
            ret += p.size();
        return ret;
    }

    Value&
    operator[](Key const& key)
    {
        return map_[partitioner(key)][key];
    }

private:
    mutable partition_map_type map_{};
};

}  // namespace ripple

#endif  // XRPL_BASICS_PARTITIONED_UNORDERED_MAP_H