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

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 key_type        = K;
    using value_type      = V;
    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 = std::allocator>
struct flat_object_set {
    using value_pointer_t = pointer_wrapper<K>;

    std::pair<value_pointer_t, bool> acquire(const K& key)
    {
        size_t hash_key      = element_hasher{}(key);
        auto [iter, is_new]  = data.emplace_with_hash(hash_key, nullptr, size_t{1});
        auto  object_pointer = const_cast<K*>(iter->first);
        auto& refcount       = iter->second;
        if (is_new) {
            object_pointer  = static_cast<K*>(mi_aligned_alloc(alignof(K), sizeof(K)));
            *object_pointer = std::move(key);
            return {make_pointer_wrapper(object_pointer), true};
        } else {
            refcount++;
            return {make_pointer_wrapper(object_pointer), false};
        }
    }

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

protected:
    using value_type     = K*;
    using element_hasher = hasher<K>;

    struct inernal_hasher {
        size_t operator()(value_type ptr) const
        {
            if (ptr == nullptr) return size_t{0};
            return element_hasher{}(*ptr);
        }
    };

    flat_hash_map_mp<value_type,
                     size_t,
                     inernal_hasher,
                     phmap::priv::hash_default_eq<value_type>,
                     Allocator<std::pair<const value_type, size_t>>>
        data{};
};

// template <typename K, template <typename> typename Allocator>
// struct flat_object_set {
//     using object_type     = K;
//     using value_pointer_t = pointer_wrapper<K>;
//     using Hasher          = hasher<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,
//                   phmap::priv::hash_default_hash<size_t>,
//                   phmap::priv::hash_default_eq<size_t>,
//                   Allocator<std::pair<const size_t, 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>
using flat_object_set_mp = detail::flat_object_set<K, std::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>;