#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include <assert.h>
#include <unistd.h>
#include <lua5.4/lua.h>
#include <lua5.4/lauxlib.h>
#include <lua5.4/lualib.h>

static void get_random(uint8_t *dst, size_t count) {
    int fd = open("/dev/urandom", O_RDONLY);
    assert(count == read(fd, dst, count));
    close(fd);
}

#if defined(ALLOC_MAGIC_BUDDY)

#include "magic_buddy/magic_buddy.h"
const size_t POOL_SIZE = 1024 * 1024 * 1024;
static struct buddy BUDDY;

static void init_allocator() {
    uint8_t magic[MAGIC_COOKIE_BYTES];
    get_random(magic, MAGIC_COOKIE_BYTES);

    uint8_t *pool = malloc(POOL_SIZE);
    init_buddy(pool, POOL_SIZE, magic, &BUDDY);
}

static void *allocator_fn(void *ud, void *ptr, size_t osize, size_t nsize) {
    // https://www.lua.org/manual/5.4/manual.html#4.6
    if (!ptr) return allocate(nsize, &BUDDY);

    if (nsize == 0) {
        liberate(ptr, osize, &BUDDY);
        return NULL;
    }

    return reallocate(ptr, nsize, osize, &BUDDY);
}

#elif defined(ALLOC_GROWING_MAGIC_BUDDY)

#include "magic_buddy/magic_buddy.h"
static size_t POOL_SIZE;
static void *POOL;
static void *END_POOL;
static struct buddy BUDDY;

static void init_allocator() {
    uint8_t magic[MAGIC_COOKIE_BYTES];
    get_random(magic, MAGIC_COOKIE_BYTES);

    POOL_SIZE = 1024;
    POOL = sbrk(POOL_SIZE);
    assert(POOL != (void *)-1);
    init_buddy(POOL, POOL_SIZE, magic, &BUDDY);
}

static int grow_pool() {
    if (sbrk(POOL_SIZE) != (void*)((uint8_t*)POOL + POOL_SIZE)) return 0;
    POOL_SIZE *= 2;
    grow_buddy(POOL, POOL_SIZE, &BUDDY);
    return 1;
}

static void *allocator_fn(void *ud, void *ptr, size_t osize, size_t nsize) {
    // https://www.lua.org/manual/5.4/manual.html#4.6
    void *result = 0;
    if (!ptr) {
        while (!(result = allocate(nsize, &BUDDY)))
            if (!grow_pool()) return 0;
        return result;
    }

    if (nsize == 0) {
        liberate(ptr, osize, &BUDDY);
        return NULL;
    }

    while (!(result = reallocate(ptr, nsize, osize, &BUDDY)))
        if (!grow_pool()) return 0;
    return result;
}

// custom malloc/free/calloc/realloc
void *malloc(size_t s) {
    if (!POOL) init_allocator();
    size_t *block = allocator_fn((void*)1, 0, 0, s + sizeof(size_t));
    *block = s + sizeof(size_t);
    return (void*)((uint8_t*)block + sizeof(size_t));
}
void free(void *ptr) {
    if (!ptr) return;
    if (!POOL) init_allocator();
    ptr = (uint8_t*)ptr - sizeof(size_t);
    size_t size = *(size_t*)ptr;
    allocator_fn((void*)1, ptr, size, 0);
}
void *realloc(void *ptr, size_t nsize) {
    if (!POOL) init_allocator();
    if (nsize == 0) {
        free(ptr);
        return 0;
    }
    ptr = (uint8_t*)ptr - sizeof(size_t);
    size_t osize = *(size_t*)ptr;
    void *res = allocator_fn((void*)1, ptr, osize, nsize + sizeof(size_t));
    *(size_t*)res = nsize + sizeof(size_t);
    return (void*)((uint8_t*)res + sizeof(size_t));
}
void *calloc(size_t cnt, size_t size) {
    if (!POOL) init_allocator();
    uint8_t *data = malloc(cnt * size);
    for (size_t i = 0; i < cnt * size; i++) data[i] = 0;
    return (void*)data;
}

#elif defined(ALLOC_SYSTEM)

static void init_allocator() { }

static void *allocator_fn(void *ud, void *ptr, size_t osize, size_t nsize) {
    // https://www.lua.org/manual/5.4/manual.html#4.6
    if (nsize == 0) {
        free(ptr);
        return NULL;
    }
    return realloc(ptr, nsize);
}

#else
#error "No allocator chosen!"
#endif

int main(int argc, char **argv) {
    init_allocator();

    // https://www.lua.org/pil/24.1.html
    char buff[256];
    int error;
    lua_State *L = lua_newstate(allocator_fn, 0);   /* opens Lua */
    luaL_openlibs(L);

    assert(argc == 2);
    luaL_dofile(L, argv[1]);
    lua_close(L);
    return 0;
}