#ifndef _BUFFERS_HPP_
#define _BUFFERS_HPP_

#include <functional>

#include "objpool.hpp"

template<uint32_t BUFSIZE>
struct MemBlock { unsigned char mem[BUFSIZE]; };

class StaticBuffer {
public:
    StaticBuffer(unsigned char* buf, uint32_t sz, std::function<void()> cb) : buffer(buf), size(sz), release_cb(cb) {};
    StaticBuffer(const StaticBuffer& sb) = delete;
    StaticBuffer(StaticBuffer&& sb) {
        buffer = sb.buffer;
        sb.buffer = nullptr;
        size = sb.size;
        sb.size = 0;
        release_cb = std::move(sb.release_cb);
    }
    ~StaticBuffer() {
        if(release_cb)
            release_cb();
    }
    unsigned char* Data() { return buffer; }
    uint32_t Capacity() { return size; }

protected:
    unsigned char* buffer = nullptr;
    uint32_t size = 0;
    std::function<void()> release_cb;
};

class BufferMgr {
public:
    StaticBuffer AllocStaticBuffer(uint32_t size) {
        if(size <= 128) {
            auto memblk = pool1.Alloc();
            return StaticBuffer(memblk->mem, 128, [this, memblk]() {
                pool1.Recycle(memblk);
            });
        } else if(size <= 1024) {
            auto memblk = pool2.Alloc();
            return StaticBuffer(memblk->mem, 1024, [this, memblk]() {
                pool2.Recycle(memblk);
            });
        } else if(size <= 4096) {
            auto memblk = pool3.Alloc();
            return StaticBuffer(memblk->mem, 4096, [this, memblk]() {
                pool3.Recycle(memblk);
            });
        } else {
            auto memblk = new unsigned char[size];
            return StaticBuffer(memblk, size, [memblk]() {
                delete[] memblk;
            });
        }
    }

protected:
    ObjectPoolNoLock<MemBlock<128>> pool1;
    ObjectPoolNoLock<MemBlock<1024>> pool2;
    ObjectPoolNoLock<MemBlock<4096>> pool3;
};

#endif