#pragma once
#include "ncx_slab.h"
#include <cstdlib>

template<int32_t alignment = 64>
class mempool {
public:
    mempool(size_t pool_size = 1024000) { // 1024KB
        const auto space = static_cast<u_char*>(malloc(pool_size));

        mem_pool_            = reinterpret_cast<ncx_slab_pool_t*>(space);
        mem_pool_->addr      = space;
        mem_pool_->min_shift = 3;
        mem_pool_->end       = space + pool_size;

        ncx_slab_init(mem_pool_);
    }
    ~mempool() {
        free(mem_pool_->addr);
    }
    [[nodiscard]] void* alloc(size_t size) const {
        return ncx_slab_alloc(mem_pool_, size);
    }

    template<typename T, typename ...Args>
    [[nodiscard]] T* alloc(Args&&... args) const {
        T* out = (T*)alloc(sizeof(T));
        new (out) T(std::forward<Args>(args)...);
        return out;
    }
    template<typename T>
    [[nodiscard]] T* alloc() const {
        T* out = (T*)alloc(sizeof(T));
        new (out) T();
        return out;
    }

    template<typename T>
    void free(T* p) const {
        p->~T();
        free((void*)p);
    }

    void free(void* p) const {
        ncx_slab_free(mem_pool_, p);
    }
    [[nodiscard]] ncx_slab_stat_t stat() const {
        ncx_slab_stat_t out;
        ncx_slab_stat(mem_pool_, &out);
        return out;
    }
private:
    ncx_slab_pool_t* mem_pool_ = nullptr;
};

template<class T>
class obj_mempool {
public:
    template<typename ...Args>
    static T* alloc(Args&&... args) {
        auto obj = pool_.template alloc<T>(std::forward<Args>(args)...);
        objs_.push_back(obj);
        return obj;
    }
    static T* alloc() {
        auto obj = pool_.template alloc<T>();
        objs_.push_back(obj);
        return obj;
    }
    static void free(T* p) {
        pool_.template free<T>(p);
        objs_.erase(std::remove(objs_.begin(), objs_.end(), p), objs_.end());
    }
    static void free_all() {
        for (auto obj : objs_) {
            pool_.template free<T>(obj);
        }
        objs_.clear();
    }
    static const std::vector<T*>& objs() {
        return objs_;
    }
    static bool has_obj(T* p) {
        return std::find(objs_.begin(), objs_.end(), p) != objs_.end();
    }
    static bool safe_free(T* p) {
        if (has_obj(p)) {
            free(p);
            return true;
        }
        return false;
    }
private:
    static mempool<alignof(T)> pool_;
    static std::vector<T*> objs_;
};