---

#pragma once

#include <unordered_map>
#include <atomic>
#include <mutex>
#include <algorithm>
#include <new>

#include <utility>

#include <limits>
#include <cstddef>

#include <bit>
#include <cassert>

namespace mm
{

#ifdef __cpp_lib_hardware_interference_size
    inline constexpr std::size_t cache_line = std::hardware_destructive_interference_size;
#else
    inline constexpr std::size_t cache_line = 64;
#endif

    class guard
    {
    protected:
        struct alloc_info {
            std::size_t packed_data;

            static alloc_info pack(std::size_t size, std::size_t alignment) noexcept {
                std::size_t log2_align = static_cast<std::size_t>(std::countr_zero(alignment));
                std::size_t packed = size | (log2_align << 58);
                return alloc_info{ packed };
            }

            std::size_t get_size() const noexcept {
                return packed_data & ((1ULL << 58) - 1);
            }

            std::size_t get_alignment() const noexcept {
                std::size_t log2_align = packed_data >> 58;
                return 1ULL << log2_align;
            }
        };

        static std::mutex& m_mm_guard_mtx() {
            alignas(mm::cache_line) static std::mutex instance;
            return instance;
        }

        alignas(mm::cache_line) inline static std::atomic<std::size_t> m_mm_bytes_allocated{0};

        static std::unordered_map<void*, alloc_info>* m_allocated()
        {
            static std::unordered_map<void*, alloc_info>* instance = new std::unordered_map<void*, alloc_info>();
            return instance;
        }

    private:
        class cleaner
        {
        public:
            ~cleaner()
            {
                std::scoped_lock lock(m_mm_guard_mtx());

                while (!m_allocated()->empty()) {
                    auto it = m_allocated()->begin();
                    void* ptr = it->first;
                    alloc_info info = it->second;
                    
                    m_allocated()->erase(it);

                    if (info.get_alignment() > __STDCPP_DEFAULT_NEW_ALIGNMENT__) {
                        ::operator delete(ptr, info.get_size(), std::align_val_t{info.get_alignment()});
                    } else {
                        ::operator delete(ptr, info.get_size());
                    }
                }

                m_mm_bytes_allocated = 0;
            }
        };

        inline static cleaner m_cleaner{};

    public:
        guard() noexcept {
            static const bool initialized = []()
            {
                std::scoped_lock lock(m_mm_guard_mtx());
                m_allocated()->reserve(1024);
                return true;
            }();
        }
        ~guard() = default;

        friend inline std::size_t allocated_bytes() noexcept;
    };

    inline std::size_t allocated_bytes() noexcept
    {
        return mm::guard::m_mm_bytes_allocated;
    }

    template <typename Ty>
    class guard_alloc : public mm::guard
    {

    public:
        using value_type = Ty;

        guard_alloc() noexcept = default;

        template <typename U>
        guard_alloc(const mm::guard_alloc<U> &) noexcept {};

        [[nodiscard]] value_type *allocate(std::size_t n)
        {
            static_assert(sizeof(value_type) != 0, "cannot allocate incomplete types");

            if (n == 0)
                return nullptr;

#define __ALLOC_LIMIT_ASSUME_COND (n <= std::numeric_limits<std::size_t>::max() / sizeof(value_type))

            assert(__ALLOC_LIMIT_ASSUME_COND && "allocation size overflow");

#if defined(__cpp_attributes) && __has_cpp_attribute(assume)
            [[assume(__ALLOC_LIMIT_ASSUME_COND)]];
#elif defined(_MSC_VER)
            __assume(__ALLOC_LIMIT_ASSUME_COND);
#elif defined(__GNUC__) || defined(__clang__)
            if (!__ALLOC_LIMIT_ASSUME_COND) 
                __builtin_unreachable();
#endif

            std::size_t alloc_bytes = n * sizeof(value_type);

            void* alloc_res = nullptr;
            if (alignof(value_type) > __STDCPP_DEFAULT_NEW_ALIGNMENT__) {
                alloc_res = ::operator new(alloc_bytes, std::align_val_t{alignof(value_type)});
            } else {
                alloc_res = ::operator new(alloc_bytes);
            }

            try {
                std::scoped_lock lock(m_mm_guard_mtx());
                m_allocated()->insert_or_assign(alloc_res, alloc_info::pack(alloc_bytes, alignof(value_type)));
            } catch(...) {
                ::operator delete(alloc_res, std::align_val_t{alignof(value_type)});
                throw;
            }

            m_mm_bytes_allocated.fetch_add(alloc_bytes, std::memory_order_relaxed);
            return static_cast<value_type *>(alloc_res);
        }

        void deallocate(value_type *p, std::size_t n) noexcept
        {
            if (!p || n == 0)
                return;

            void *target_ptr = static_cast<void *>(p);
            alloc_info info{};
            bool was_erased = false;
        
            {
                std::scoped_lock lock(m_mm_guard_mtx());
                auto it = m_allocated()->find(target_ptr);
                if (it != m_allocated()->end()) {
                    info = it->second;
                    m_allocated()->erase(it);
                    was_erased = true;
                }
            }

            if(was_erased) {
                std::size_t dealloc_bytes = info.get_size();
                std::size_t alignment = info.get_alignment();

                if (alignment > __STDCPP_DEFAULT_NEW_ALIGNMENT__) {
                    ::operator delete(target_ptr, dealloc_bytes, std::align_val_t{alignment});
                } else {
                    ::operator delete(target_ptr, dealloc_bytes);
                }
                m_mm_bytes_allocated.fetch_sub(dealloc_bytes, std::memory_order_relaxed);
            }
        }

        template <typename U>
        friend bool operator==(const guard_alloc&, const guard_alloc<U>&) noexcept { return true; }

        template <typename U>
        friend bool operator!=(const guard_alloc&, const guard_alloc<U>&) noexcept { return false; }
    };

    namespace ga
    {
        template<typename MakeTy, typename... Args>
        [[nodiscard]] MakeTy* make(Args&&... args) {
            auto alloc = guard_alloc<MakeTy>();
            MakeTy *ptr = alloc.allocate(1);
            if(ptr) 
                ::new (static_cast<void *>(ptr)) MakeTy(std::forward<Args>(args)...);

            return ptr;
        }

        template <typename DropTy>
        void drop(DropTy *&ptr) noexcept
        {
            if (!ptr)
                return;

            ptr->~DropTy();
            
            auto alloc = guard_alloc<DropTy>();
            alloc.deallocate(ptr, 1);
            ptr = nullptr;
        }

    } // namespace ga

} // namespace mm