#include <container/hashmap.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;
            value.object_pointer = mi_aligned_alloc(alignof(V), sizeof(V));
            if constexpr (std::is_same_v<V, K>)
                *value.object_pointer = key;
            else
                *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, V>>> data{};
};
} // 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>;