//
// immer: immutable data structures for C++
// Copyright (C) 2016, 2017, 2018 Juan Pedro Bolivar Puente
//
// This software is distributed under the Boost Software License, Version 1.0.
// See accompanying file LICENSE or copy at http://boost.org/LICENSE_1_0.txt
//

#pragma once

#include <immer/config.hpp>
#include <immer/detail/hamts/node.hpp>

#include <algorithm>

namespace immer {
namespace detail {
namespace hamts {

template <typename T,
          typename Hash,
          typename Equal,
          typename MemoryPolicy,
          bits_t B>
struct champ
{
    static constexpr auto bits = B;

    using node_t   = node<T, Hash, Equal, MemoryPolicy, B>;
    using bitmap_t = typename get_bitmap_type<B>::type;

    static_assert(branches<B> <= sizeof(bitmap_t) * 8, "");

    node_t* root;
    size_t size;

    static const champ& empty()
    {
        static const champ empty_{
            node_t::make_inner_n(0),
            0,
        };
        return empty_;
    }

    champ(node_t* r, size_t sz)
        : root{r}
        , size{sz}
    {}

    champ(const champ& other)
        : champ{other.root, other.size}
    {
        inc();
    }

    champ(champ&& other)
        : champ{empty()}
    {
        swap(*this, other);
    }

    champ& operator=(const champ& other)
    {
        auto next = other;
        swap(*this, next);
        return *this;
    }

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

    friend void swap(champ& x, champ& y)
    {
        using std::swap;
        swap(x.root, y.root);
        swap(x.size, y.size);
    }

    ~champ() { dec(); }

    void inc() const { root->inc(); }

    void dec() const
    {
        if (root->dec())
            node_t::delete_deep(root, 0);
    }

    template <typename Fn>
    void for_each_chunk(Fn&& fn) const
    {
        for_each_chunk_traversal(root, 0, fn);
    }

    template <typename Fn>
    void for_each_chunk_traversal(node_t* node, count_t depth, Fn&& fn) const
    {
        if (depth < max_depth<B>) {
            auto datamap = node->datamap();
            if (datamap)
                fn(node->values(), node->values() + popcount(datamap));
            auto nodemap = node->nodemap();
            if (nodemap) {
                auto fst = node->children();
                auto lst = fst + popcount(nodemap);
                for (; fst != lst; ++fst)
                    for_each_chunk_traversal(*fst, depth + 1, fn);
            }
        } else {
            fn(node->collisions(),
               node->collisions() + node->collision_count());
        }
    }

    template <typename Project, typename Default, typename K>
    decltype(auto) get(const K& k) const
    {
        auto node = root;
        auto hash = Hash{}(k);
        for (auto i = count_t{}; i < max_depth<B>; ++i) {
            auto bit = bitmap_t{1u} << (hash & mask<B>);
            if (node->nodemap() & bit) {
                auto offset = popcount(node->nodemap() & (bit - 1));
                node        = node->children()[offset];
                hash        = hash >> B;
            } else if (node->datamap() & bit) {
                auto offset = popcount(node->datamap() & (bit - 1));
                auto val    = node->values() + offset;
                if (Equal{}(*val, k))
                    return Project{}(*val);
                else
                    return Default{}();
            } else {
                return Default{}();
            }
        }
        auto fst = node->collisions();
        auto lst = fst + node->collision_count();
        for (; fst != lst; ++fst)
            if (Equal{}(*fst, k))
                return Project{}(*fst);
        return Default{}();
    }

    std::pair<node_t*, bool>
    do_add(node_t* node, T v, hash_t hash, shift_t shift) const
    {
        if (shift == max_shift<B>) {
            auto fst = node->collisions();
            auto lst = fst + node->collision_count();
            for (; fst != lst; ++fst)
                if (Equal{}(*fst, v))
                    return {
                        node_t::copy_collision_replace(node, fst, std::move(v)),
                        false};
            return {node_t::copy_collision_insert(node, std::move(v)), true};
        } else {
            auto idx = (hash & (mask<B> << shift)) >> shift;
            auto bit = bitmap_t{1u} << idx;
            if (node->nodemap() & bit) {
                auto offset = popcount(node->nodemap() & (bit - 1));
                auto result = do_add(
                    node->children()[offset], std::move(v), hash, shift + B);
                try {
                    result.first =
                        node_t::copy_inner_replace(node, offset, result.first);
                    return result;
                } catch (...) {
                    node_t::delete_deep_shift(result.first, shift + B);
                    throw;
                }
            } else if (node->datamap() & bit) {
                auto offset = popcount(node->datamap() & (bit - 1));
                auto val    = node->values() + offset;
                if (Equal{}(*val, v))
                    return {node_t::copy_inner_replace_value(
                                node, offset, std::move(v)),
                            false};
                else {
                    auto child = node_t::make_merged(
                        shift + B, std::move(v), hash, *val, Hash{}(*val));
                    try {
                        return {node_t::copy_inner_replace_merged(
                                    node, bit, offset, child),
                                true};
                    } catch (...) {
                        node_t::delete_deep_shift(child, shift + B);
                        throw;
                    }
                }
            } else {
                return {
                    node_t::copy_inner_insert_value(node, bit, std::move(v)),
                    true};
            }
        }
    }

    champ add(T v) const
    {
        auto hash     = Hash{}(v);
        auto res      = do_add(root, std::move(v), hash, 0);
        auto new_size = size + (res.second ? 1 : 0);
        return {res.first, new_size};
    }

    template <typename Project,
              typename Default,
              typename Combine,
              typename K,
              typename Fn>
    std::pair<node_t*, bool>
    do_update(node_t* node, K&& k, Fn&& fn, hash_t hash, shift_t shift) const
    {
        if (shift == max_shift<B>) {
            auto fst = node->collisions();
            auto lst = fst + node->collision_count();
            for (; fst != lst; ++fst)
                if (Equal{}(*fst, k))
                    return {
                        node_t::copy_collision_replace(
                            node,
                            fst,
                            Combine{}(std::forward<K>(k),
                                      std::forward<Fn>(fn)(Project{}(*fst)))),
                        false};
            return {node_t::copy_collision_insert(
                        node,
                        Combine{}(std::forward<K>(k),
                                  std::forward<Fn>(fn)(Default{}()))),
                    true};
        } else {
            auto idx = (hash & (mask<B> << shift)) >> shift;
            auto bit = bitmap_t{1u} << idx;
            if (node->nodemap() & bit) {
                auto offset = popcount(node->nodemap() & (bit - 1));
                auto result = do_update<Project, Default, Combine>(
                    node->children()[offset],
                    k,
                    std::forward<Fn>(fn),
                    hash,
                    shift + B);
                try {
                    result.first =
                        node_t::copy_inner_replace(node, offset, result.first);
                    return result;
                } catch (...) {
                    node_t::delete_deep_shift(result.first, shift + B);
                    throw;
                }
            } else if (node->datamap() & bit) {
                auto offset = popcount(node->datamap() & (bit - 1));
                auto val    = node->values() + offset;
                if (Equal{}(*val, k))
                    return {
                        node_t::copy_inner_replace_value(
                            node,
                            offset,
                            Combine{}(std::forward<K>(k),
                                      std::forward<Fn>(fn)(Project{}(*val)))),
                        false};
                else {
                    auto child = node_t::make_merged(
                        shift + B,
                        Combine{}(std::forward<K>(k),
                                  std::forward<Fn>(fn)(Default{}())),
                        hash,
                        *val,
                        Hash{}(*val));
                    try {
                        return {node_t::copy_inner_replace_merged(
                                    node, bit, offset, child),
                                true};
                    } catch (...) {
                        node_t::delete_deep_shift(child, shift + B);
                        throw;
                    }
                }
            } else {
                return {node_t::copy_inner_insert_value(
                            node,
                            bit,
                            Combine{}(std::forward<K>(k),
                                      std::forward<Fn>(fn)(Default{}()))),
                        true};
            }
        }
    }

    template <typename Project,
              typename Default,
              typename Combine,
              typename K,
              typename Fn>
    champ update(const K& k, Fn&& fn) const
    {
        auto hash = Hash{}(k);
        auto res  = do_update<Project, Default, Combine>(
            root, k, std::forward<Fn>(fn), hash, 0);
        auto new_size = size + (res.second ? 1 : 0);
        return {res.first, new_size};
    }

    // basically:
    //      variant<monostate_t, T*, node_t*>
    // boo bad we are not using... C++17 :'(
    struct sub_result
    {
        enum kind_t
        {
            nothing,
            singleton,
            tree
        };

        union data_t
        {
            T* singleton;
            node_t* tree;
        };

        kind_t kind;
        data_t data;

        sub_result()
            : kind{nothing} {};
        sub_result(T* x)
            : kind{singleton}
        {
            data.singleton = x;
        };
        sub_result(node_t* x)
            : kind{tree}
        {
            data.tree = x;
        };
    };

    template <typename K>
    sub_result
    do_sub(node_t* node, const K& k, hash_t hash, shift_t shift) const
    {
        if (shift == max_shift<B>) {
            auto fst = node->collisions();
            auto lst = fst + node->collision_count();
            for (auto cur = fst; cur != lst; ++cur)
                if (Equal{}(*cur, k))
                    return node->collision_count() > 2
                               ? node_t::copy_collision_remove(node, cur)
                               : sub_result{fst + (cur == fst)};
            return {};
        } else {
            auto idx = (hash & (mask<B> << shift)) >> shift;
            auto bit = bitmap_t{1u} << idx;
            if (node->nodemap() & bit) {
                auto offset = popcount(node->nodemap() & (bit - 1));
                auto result =
                    do_sub(node->children()[offset], k, hash, shift + B);
                switch (result.kind) {
                case sub_result::nothing:
                    return {};
                case sub_result::singleton:
                    return node->datamap() == 0 &&
                                   popcount(node->nodemap()) == 1 && shift > 0
                               ? result
                               : node_t::copy_inner_replace_inline(
                                     node, bit, offset, *result.data.singleton);
                case sub_result::tree:
                    try {
                        return node_t::copy_inner_replace(
                            node, offset, result.data.tree);
                    } catch (...) {
                        node_t::delete_deep_shift(result.data.tree, shift + B);
                        throw;
                    }
                }
            } else if (node->datamap() & bit) {
                auto offset = popcount(node->datamap() & (bit - 1));
                auto val    = node->values() + offset;
                if (Equal{}(*val, k)) {
                    auto nv = popcount(node->datamap());
                    if (node->nodemap() || nv > 2)
                        return node_t::copy_inner_remove_value(
                            node, bit, offset);
                    else if (nv == 2) {
                        return shift > 0 ? sub_result{node->values() + !offset}
                                         : node_t::make_inner_n(
                                               0,
                                               node->datamap() & ~bit,
                                               node->values()[!offset]);
                    } else {
                        assert(shift == 0);
                        return empty().root->inc();
                    }
                }
            }
            return {};
        }
    }

    template <typename K>
    champ sub(const K& k) const
    {
        auto hash = Hash{}(k);
        auto res  = do_sub(root, k, hash, 0);
        switch (res.kind) {
        case sub_result::nothing:
            return *this;
        case sub_result::tree:
            return {res.data.tree, size - 1};
        default:
            IMMER_UNREACHABLE;
        }
    }

    template <typename Eq = Equal>
    bool equals(const champ& other) const
    {
        return size == other.size && equals_tree<Eq>(root, other.root, 0);
    }

    template <typename Eq>
    static bool equals_tree(const node_t* a, const node_t* b, count_t depth)
    {
        if (a == b)
            return true;
        else if (depth == max_depth<B>) {
            auto nv = a->collision_count();
            return nv == b->collision_count() &&
                   equals_collisions<Eq>(a->collisions(), b->collisions(), nv);
        } else {
            if (a->nodemap() != b->nodemap() || a->datamap() != b->datamap())
                return false;
            auto n = popcount(a->nodemap());
            for (auto i = count_t{}; i < n; ++i)
                if (!equals_tree<Eq>(
                        a->children()[i], b->children()[i], depth + 1))
                    return false;
            auto nv = popcount(a->datamap());
            return !nv || equals_values<Eq>(a->values(), b->values(), nv);
        }
    }

    template <typename Eq>
    static bool equals_values(const T* a, const T* b, count_t n)
    {
        return std::equal(a, a + n, b, Eq{});
    }

    template <typename Eq>
    static bool equals_collisions(const T* a, const T* b, count_t n)
    {
        auto ae = a + n;
        auto be = b + n;
        for (; a != ae; ++a) {
            for (auto fst = b; fst != be; ++fst)
                if (Eq{}(*a, *fst))
                    goto good;
            return false;
        good:
            continue;
        }
        return true;
    }
};

} // namespace hamts
} // namespace detail
} // namespace immer