#ifndef APP_ECOM_IM_MEMPOOL_SRC_MEMPOOL_H
#define APP_ECOM_IM_MEMPOOL_SRC_MEMPOOL_H

#include <execinfo.h>
#include <sstream>
#include <pthread.h>
#include <new>
#include <butil/atomicops.h>
#include <butil/logging.h>
#include <iostream>

namespace im {
namespace fugue {

namespace lockfree {

template <class T>
class PushOnlyStack {
public:
    PushOnlyStack() {
        _head.store(NULL, butil::memory_order_relaxed);
    }

    void push(T* node) {
        T* head = _head.load(butil::memory_order_relaxed);
        node->next = head;
        while (!_head.compare_exchange_weak(
                       head, node, butil::memory_order_relaxed)) {
            node->next = head;
        }
    }

    T* release() {
        return _head.exchange(NULL, butil::memory_order_relaxed);
    }

private:
    butil::atomic<T*> _head;
};

template <class T>
struct FreeListNode {
    uint64_t id;
    uint64_t next;
    T data;
};

template <class T, int CAP>
class FreeList {
public:
    typedef FreeListNode<T> Node;
    static const uint64_t EMPTY = 0xFFFFFFFFFFFFFFFF;

    T* get() {
        uint64_t head = _head.load(butil::memory_order_acquire);
        if (head == EMPTY) {
            return new_node();
        }

        Node* node = address(head);
        while (!_head.compare_exchange_weak(
                       head, node->next, butil::memory_order_acquire)) {
            if (head == EMPTY) {
                return new_node();
            }
            node = address(head);
        }
        return &node->data;
    }

    void put(T* value) {
        Node* node = container_of(value, Node, data);

        uint64_t head = _head.load(butil::memory_order_acquire);
        // add version
        node->id += (1UL << 32);
        node->next = head;

        // NOTE: we MUST use a temp var *head* to call compare_exchange_weak
        // because Boost.Atomic will update the *expected* even success
        // std::atomic do not have this limitation
        while (!_head.compare_exchange_weak(
                       head, node->id, butil::memory_order_release)) {
            node->next = head;
        }
    }

    template <class F>
    void unsafe_foreach() {
        uint32_t used_blk_cnt = _slot_index.load(butil::memory_order_relaxed);
        for (uint32_t i = 0; i < used_blk_cnt; ++i) {
            F()(&_node[i]->data);
        }
    }

    uint32_t real_used_size() const {
        uint32_t used_blk_cnt = _slot_index.load(butil::memory_order_relaxed);
        uint64_t used_bytes = 0;
        for (uint32_t i = 0; i < used_blk_cnt; ++i) {
            used_bytes += _node[i]->data.offset;
        }
        return used_bytes >> 10;
    }

    uint32_t allocate_blocks() const {
        return _slot_index.load(butil::memory_order_relaxed);
    }

    uint32_t free_blocks() const {
        uint64_t head = _head.load(butil::memory_order_relaxed);
        uint32_t size = 0;
        while (head != FreeList::EMPTY) {
            const Node* head_ptr = address(head);
            head = head_ptr->next;
            ++size;
        }
        return size;
    }

    void reset() {
        _head.store(FreeList::EMPTY, butil::memory_order_relaxed);
        _slot_index.store(0, butil::memory_order_relaxed);
    }

    FreeList() {
        for (int i = 0; i < CAP; ++i) {
            _node[i] = NULL;
        }
        reset();
    }

private:
    uint32_t slot(uint64_t id) const {
        return static_cast<uint32_t>(id);
    }

    T* new_node() {
        uint32_t index = _slot_index.fetch_add(1, butil::memory_order_relaxed);
        if (index >= CAP) {
            return NULL;
        }

        if (_node[index] != NULL) {
            return &(_node[index]->data);
        }

        Node* node = (Node*)malloc(sizeof(Node));
        new (node) Node;

        node->id = index;
        _node[index] = node;

        return &node->data;
    }

    Node* address(uint64_t id) {
        return _node[slot(id)];
    }

    const Node* address(uint64_t id) const {
        return _node[slot(id)];
    }

    butil::atomic<uint64_t> _head;
    butil::atomic<uint32_t> _slot_index;
    Node* _node[CAP];
};

}

namespace memory {

struct Block {
    static const int BLOCK_SIZE = 2 * 1024 * 1024;
    char data[BLOCK_SIZE];
};

class GlobalBlockFreeList {
public:
    static const int MAX_BLOCK_COUNT = 32 * 1024;
    typedef lockfree::FreeList<Block, MAX_BLOCK_COUNT> type;
    static type* instance() {
        static type singleton;
        return &singleton;
    }
};

struct BlockReference {
    BlockReference() : offset(0), block(NULL) {
        // do nothing
    }

    void reset() {
        offset = 0;
        block = NULL;
    }

    uint32_t offset;
    Block* block;
};

class Region {
public:
    struct GlobalPut {
        void operator()(BlockReference* block_ref) {
            if (block_ref->block != NULL) {
                GlobalBlockFreeList::instance()->put(block_ref->block);
            }
            block_ref->reset();
        }
    };

    struct BigNode {
        BigNode* next;
        char data[0];
    };

    ~Region() {
        reset();
        delete [] _big_mem_start;
        _big_mem_start = NULL;
    }

    char const* debug_str() const {
        uint32_t alloc_blocks = _free_blocks.allocate_blocks();
        uint32_t free_blocks = _free_blocks.free_blocks();
        uint32_t used_mem_mb = _free_blocks.real_used_size();
        uint32_t big_buf_size = _big_mem_size.load(butil::memory_order_relaxed);
        uint32_t big_buf_count = _big_mem_count.load(butil::memory_order_relaxed);
        uint32_t mlc_mem_size = _mlc_mem_size.load(butil::memory_order_relaxed);
        uint32_t mlc_mem_count = _mlc_mem_count.load(butil::memory_order_relaxed);

        std::ostringstream oss;
        oss << "[alloc_blks:" << alloc_blocks << ",free_blks:" << free_blocks
            << ",used_mem_kb:" << used_mem_mb << ",big_mem_kb:" << (big_buf_size >> 10)
            << ",big_buf_cnt:" << big_buf_count << ",mlc_mem_kb:" << (mlc_mem_size >> 10)
            << ",mlc_cnt:" << mlc_mem_count << "]";

        return oss.str().c_str();
    }

    Region();

    void init();

    void reset();

    BlockReference* get();

    void* malloc(size_t size);

    void put(BlockReference* block);

    static const int MAX_BLOCK_COUNT = 1024;
    static const int BIG_MEM_THRESHOLD = 256 * 1024;
    static const int MLC_MEM_THRESHOLD = 4 * 1024 * 1024;
    static const int COUNTER_SIZE = MLC_MEM_THRESHOLD / BIG_MEM_THRESHOLD + 1;

private:
    lockfree::FreeList<BlockReference, MAX_BLOCK_COUNT> _free_blocks;
    lockfree::PushOnlyStack<BigNode> _big_nodes;

    butil::atomic<uint32_t> _big_mem_size;
    butil::atomic<uint32_t> _big_mem_count;

    char* _big_mem_start;
    uint32_t _big_mem_capacity;

    butil::atomic<uint32_t> _mlc_mem_size;
    butil::atomic<uint32_t> _mlc_mem_count;
};

}

}

class Mempool {
public:
    void* malloc(size_t size) {
        size = _align(size);
        if (size <= _free_size) {
            void* p = _free_cursor;
            _free_size -= size;
            _free_cursor += size;
            return p;
        }

        return malloc_from_region(size);
    }

    void free(void* p, size_t size) {
        if (size >= fugue::memory::Region::BIG_MEM_THRESHOLD) {
            return;
        }

        if (_free_cursor - size == static_cast<char*>(p)) {
            size_t down_aligned = _down_align(size);
            _free_cursor -= down_aligned;
            _free_size += down_aligned;
        }
    }

    void* realloc(void* old_data, size_t old_size, size_t new_size) {
        if (old_size >= new_size) {
            return old_data;
        }

        size_t required = new_size - old_size;
        if (_free_cursor == static_cast<char*>(old_data) + old_size) {
            if (_free_size >= required) {
                _free_cursor += required;
                _free_size -= required;
                return old_data;
            } else {
                _free_cursor = static_cast<char*>(old_data);
                _free_size += old_size;
            }
        }

        void* p = this->malloc_from_region(new_size);
        if (p != NULL) {
            memcpy(p, old_data, old_size);
            return p;
        }

        return NULL;
    }

    Mempool(fugue::memory::Region* blocks) : _free_size(0)
             , _free_cursor(NULL)
             , _blocks(blocks) {
        _block = NULL;
    }

    ~Mempool() {
        release_block();
    }

    void release_block() {
        if (_block) {
            _block->offset = fugue::memory::Block::BLOCK_SIZE - _free_size;
            _blocks->put(_block);
        }

        _free_size = 0;
        _free_cursor = NULL;
        _block = NULL;
    }

private:
    void* malloc_from_region(size_t size) {
        if (size >= fugue::memory::Region::BIG_MEM_THRESHOLD) {
            return _blocks->malloc(size);
        }

        while (true) {
            fugue::memory::BlockReference* block = _blocks->get();
            if (block == NULL) {
                return NULL;
            }

            uint32_t free_size = fugue::memory::Block::BLOCK_SIZE - block->offset;
            if (size <= free_size) {
                if (_block) {
                    _block->offset = fugue::memory::Block::BLOCK_SIZE - _free_size;
                }

                char* p = block->block->data + block->offset;
                _free_size = free_size - size;
                _free_cursor = p + size;
                _block = block;
                return p;
            }
        }
        return _blocks->malloc(size);
    }

    static const int ALIGN_SIZE = sizeof(void*);

    inline size_t _align(size_t size) const {
        return (size + (ALIGN_SIZE - 1)) & ~(ALIGN_SIZE - 1);
    }

    inline size_t _down_align(size_t size) const {
        return size & ~(ALIGN_SIZE - 1);
    }

    size_t _free_size;
    char* _free_cursor;

    fugue::memory::Region* _blocks;
    fugue::memory::BlockReference* _block;
};

extern __thread Mempool* g_mempool;
class mempool {
public:
    virtual void * malloc (size_t size) = 0;
    virtual void free (void *p, size_t size) = 0;
    inline virtual ~mempool(){}
};

class GlobalMempool : public mempool {
public:
    GlobalMempool() {
        // do nothing;
    }

    virtual ~GlobalMempool() {
        // do nothing;
    }

    static GlobalMempool* instance() {
        static GlobalMempool singleton;
        return &singleton;
    }

    void reset(Mempool* mempool) {
        g_mempool = mempool;
    }

    void* malloc(size_t size) {
        return g_mempool->malloc(size);
    }

    void* realloc(void* old_data, size_t old_size, size_t new_size) {
        return g_mempool->realloc(old_data, old_size, new_size);
    }

    void free(void* p, size_t s) {
        g_mempool->free(p, s);
    }

    void clear() {
        g_mempool->release_block();
    }

    Mempool* get() {
        return g_mempool;
    }

};

class MempoolGuard {
public:
    MempoolGuard(fugue::memory::Region* region) : _mempool(region) {
        acquire();
    }

    void acquire() {
        _saved_mempool = g_mempool;
        g_mempool = &_mempool;
    }

    void release() {
        _mempool.release_block();
        g_mempool = _saved_mempool;
    }

    ~MempoolGuard() {
        release();
    }

private:
    Mempool _mempool;
    Mempool* _saved_mempool;
};

inline std::string print_trace() {
    const static int BT_BUF_SIZE = 400;
    std::stringstream debug_stream;

    void* buffer[BT_BUF_SIZE];
    int nptrs = backtrace(buffer, BT_BUF_SIZE);
    char** strings = backtrace_symbols(buffer, nptrs);

    for (int j = 0; j < nptrs; j++) {
        debug_stream << strings[j] << "\t";
    }

    return debug_stream.str();
}

}
#endif