#pragma once
#include <vector>
#include <stdexcept>
#ifndef _WIN32
#  include <pthread.h>
#endif

/*
 * Static (singleton) threaded resources allocated by StaticMemoryPool<T> and ObjectAllocator are freed:
 * - On UNIX: by setting a destroy callback to pthread library for threaded instance singleton key
 * - On Windows: from DllMain located in lib.cc. In this case we track list of created StaticMemoryPool objects in vector object.
 */

namespace panda { namespace lib {

struct MemoryPool {
    MemoryPool (size_t blocksize) : first_free(NULL) {
        this->blocksize = blocksize > 8 ? blocksize : 8;
    }

    void* alloc () {
        if (!first_free) grow();
        void* ret = first_free;
        first_free = *((void**)ret);
        return ret;
    }

    void dealloc (void* elem) {
        //assert(is_mine(elem)); // protection for debugging, normally you MUST NEVER pass a pointer that wasn't created via current mempool
        *((void**)elem) = first_free;
        first_free = elem;
    }

    ~MemoryPool ();

    template <int T> friend struct StaticMemoryPool;

#ifdef _WIN32
    static void threaded_instances_free ();
#endif

private:
    struct Pool {
        char*  list;
        size_t size;
        size_t len;
    };
    size_t            blocksize;
    std::vector<Pool> pools;
    void*             first_free;

    void grow    ();
    bool is_mine (void* elem);

#ifdef _WIN32
    static thread_local std::vector<MemoryPool*>* threaded_static_pools;
    static void _thr_instance_add (MemoryPool*);
#endif
};

template <int BLOCKSIZE> struct StaticMemoryPool {

    static inline MemoryPool* instance () {
        static MemoryPool _instance(BLOCKSIZE);
        return &_instance;
    }

    static inline MemoryPool* threaded_instance () {
        if (!_thr_instance) _thr_init();
        return _thr_instance;
    }

private:
    static thread_local MemoryPool* _thr_instance;
#ifndef _WIN32
    static pthread_key_t  _thr_key;
    static pthread_once_t _thr_once;
#endif

    StaticMemoryPool () {}

#ifdef _WIN32
    static void _thr_init () {
        _thr_instance = new MemoryPool(BLOCKSIZE);
        MemoryPool::_thr_instance_add(_thr_instance);
    }
#else
    static void _thr_once_init () {
        int err = pthread_key_create(&_thr_key, _thr_delete);
        if (err) { throw std::logic_error("[MemoryPool._thr_once_init]: pthread_key_create error"); }
    }

    static void _thr_init () {
        _thr_instance = new MemoryPool(BLOCKSIZE);
        pthread_once(&_thr_once, _thr_once_init);
        int err = pthread_setspecific(_thr_key, _thr_instance);
        if (err) throw std::logic_error("[MemoryPool.threaded_instance]: pthread_setspecific error");
    }

    static void _thr_delete (void* pool) {
        delete static_cast<MemoryPool*>(pool);
    }
#endif
};

template <int T> thread_local MemoryPool* StaticMemoryPool<T>::_thr_instance;
#ifndef _WIN32
template <int T> pthread_key_t StaticMemoryPool<T>::_thr_key;
template <int T> pthread_once_t StaticMemoryPool<T>::_thr_once = PTHREAD_ONCE_INIT;
#endif

template <> struct StaticMemoryPool<7> : StaticMemoryPool<8> {};
template <> struct StaticMemoryPool<6> : StaticMemoryPool<8> {};
template <> struct StaticMemoryPool<5> : StaticMemoryPool<8> {};
template <> struct StaticMemoryPool<4> : StaticMemoryPool<8> {};
template <> struct StaticMemoryPool<3> : StaticMemoryPool<8> {};
template <> struct StaticMemoryPool<2> : StaticMemoryPool<8> {};
template <> struct StaticMemoryPool<1> : StaticMemoryPool<8> {};

struct ObjectAllocator {

    static inline ObjectAllocator* instance () {
        static ObjectAllocator _instance;
        return &_instance;
    }

    static inline ObjectAllocator* threaded_instance () {
        if (!_thr_instance) _thr_init();
        return _thr_instance;
    }

#ifdef _WIN32
    static void threaded_instance_free ();
#endif

    ObjectAllocator ();

    void* alloc (size_t size) {
        if (size == 0) return NULL;
        MemoryPool* pool;
        if (size <= 1024) {
            pool = small_pools[(size-1)>>2];
            if (!pool) pool = small_pools[(size-1)>>2] = new MemoryPool((((size-1)>>2) + 1)<<2);
        }
        else if (size <= 16384) {
            pool = medium_pools[(size-1)>>6];
            if (!pool) pool = medium_pools[(size-1)>>6] = new MemoryPool((((size-1)>>6) + 1)<<6);
        } else {
            if (size > 262144) throw std::invalid_argument("ObjectAllocator: object size cannot exceed 256k");
            pool = big_pools[(size-1)>>10];
            if (!pool) pool = big_pools[(size-1)>>10] = new MemoryPool((((size-1)>>10) + 1)<<10);
        }
        return pool->alloc();
    }

    void dealloc (void* ptr, size_t size) {
        if (ptr == NULL || size == 0) return;
        MemoryPool* pool;
        if (size <= 1024) pool = small_pools[(size-1)>>2];
        else if (size <= 16384) pool = medium_pools[(size-1)>>6];
        else {
            if (size > 262144) throw std::invalid_argument("ObjectAllocator: object size cannot exceed 256k");
            pool = big_pools[(size-1)>>10];
        }
        pool->dealloc(ptr);
    }

    ~ObjectAllocator ();

private:
    static const int POOLS_CNT = 256;
    static thread_local ObjectAllocator* _thr_instance;
#ifndef _WIN32
    static pthread_key_t  _thr_key;
    static pthread_once_t _thr_once;
#endif
    MemoryPool** small_pools;
    MemoryPool** medium_pools;
    MemoryPool** big_pools;

    static void _thr_init ();
#ifndef _WIN32
    static void _thr_once_init ();
    static void _thr_delete    (void* allocator);
#endif
};

template <class TARGET, bool THREADED = true> struct AllocatedObject {
    static void* operator new (size_t size) {
        if (size == sizeof(TARGET)) return StaticMemoryPool<sizeof(TARGET)>::threaded_instance()->alloc();
        return ObjectAllocator::threaded_instance()->alloc(size);
    }

    static void operator delete (void* p, size_t size) {
        if (size == sizeof(TARGET)) StaticMemoryPool<sizeof(TARGET)>::threaded_instance()->dealloc(p);
        else ObjectAllocator::threaded_instance()->dealloc(p, size);
    }
};

template <class TARGET> struct AllocatedObject<TARGET, false> {
    static void* operator new (size_t size) {
        if (size == sizeof(TARGET)) return StaticMemoryPool<sizeof(TARGET)>::instance()->alloc();
        return ObjectAllocator::instance()->alloc(size);
    }

    static void operator delete (void* p, size_t size) {
        if (size == sizeof(TARGET)) StaticMemoryPool<sizeof(TARGET)>::instance()->dealloc(p);
        else ObjectAllocator::instance()->dealloc(p, size);
    }
};

}}
