#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 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 *)-1) 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;
}

#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;
}