#include <container/hashmap.hpp>
#include <container/hive.hpp>
#include <utils/pointer_wrapper.hpp>

namespace detail
{
template <typename K>
struct hasher {
    size_t operator()(const K& k) const;
};

template <typename K, typename Tag>
struct tagged_hasher {
    size_t operator()(const K& k) const;
};

template <typename K>
struct eq_compare {
    bool operator()(const K& lhs, const K& rhs) const { return lhs == rhs; }
};

template <typename K, typename V>
struct default_elem_ctor {
    V operator()(const K& k) const;
};

template <typename K, typename V, typename Hasher = hasher<K>, template <typename> typename Allocator = std::allocator>
struct flat_object_map {
    using value_pointer_t = pointer_wrapper<V>;

    std::pair<value_pointer_t, bool> acquire(const K& key)
    {
        auto [iter, is_new] = data.try_emplace(key, object_with_refcount{});
        auto& value         = iter->second;
        if (is_new) {
            value.refcount        = 1;
            V* new_value          = static_cast<V*>(mi_aligned_alloc(alignof(V), sizeof(V)));
            value.object_pointer  = make_pointer_wrapper(new_value);
            *value.object_pointer = std::move(default_elem_ctor<K, V>{}(key));
            return {value.object_pointer, true};
        } else {
            value.refcount++;
            return {value.object_pointer, false};
        }
    }

    void release(const K& key)
    {
        if (auto iter = data.find(key); iter == data.end()) {
            throw std::runtime_error("Key not found in refcount map.");
        } else if (--iter->second.refcount == 0) {
            data.erase(iter);
        }
    }

protected:
    struct object_with_refcount {
        size_t          refcount{};
        value_pointer_t object_pointer{};
    };

    flat_hash_map<K, object_with_refcount, Hasher, eq_compare<K>, Allocator<std::pair<const K, object_with_refcount>>> data{};
};

template <typename K, template <typename> typename Allocator>
struct flat_object_map<K, K, hasher<size_t>, Allocator> {
    using value_pointer_t = pointer_wrapper<K>;

    std::pair<value_pointer_t, bool> acquire(const K& key)
    {
        size_t hash_key               = Hasher()(key);
        auto [iter, is_new]           = refcount_data_map.try_emplace(hash_key, object_with_refcount{});
        auto& [refcount, object_iter] = iter->second;
        if (is_new) {
            refcount    = 1;
            object_iter = data.emplace(key);
            return {make_pointer_wrapper(object_iter.operator->()), true};
        } else {
            refcount++;
            return {make_pointer_wrapper(object_iter.operator->()), false};
        }
    }

    void release(const K& key)
    {
        size_t hash_key = Hasher()(key);
        if (auto iter = refcount_data_map.find(hash_key); iter == refcount_data_map.end()) {
            throw std::runtime_error("Key not found in refcount map.");
        } else if (--iter->second.refcount == 0) {
            data.erase(iter->second.object_iter);
            refcount_data_map.erase(iter);
        }
    }

protected:
    hive<K, Allocator<K>> data{};
    using iterator = typename decltype(data)::iterator;

    struct object_with_refcount {
        size_t   refcount{};
        iterator object_iter{};
    };

    flat_hash_map<size_t,
                  object_with_refcount,
                  hasher<size_t>,
                  eq_compare<K>,
                  Allocator<std::pair<const K, object_with_refcount>>>
        refcount_data_map{};
};
} // namespace detail

template <typename K, typename V>
using flat_object_map_mp = detail::flat_object_map<K, V, detail::hasher<K>, mi_stl_allocator>;

template <typename K, typename V, typename Tag>
using tagged_flat_object_map_mp = detail::flat_object_map<K, V, detail::tagged_hasher<K, Tag>, mi_stl_allocator>;