#include <sys/mman.h>

#include <assert.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>

#include "arena.h"

#if __STDC_VERSION__ < 201112L /* below C11 */
typedef intptr_t max_align_t;
#define thread_local __thread
#endif

static void arena_defer_cleanup(arena_t *arena);

arena_t *arena_new(size_t sz)
{
    arena_t *arena;
    size_t msz = sz + sizeof(arena_t);

    if (msz < sysconf(_SC_PAGE_SIZE))
    {
        arena = malloc(msz);
    }
    else
    {
        arena = mmap(NULL, msz, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
        if (arena == MAP_FAILED) return false;
    }

    arena->a_sz = sz;
    arena->a_pos = 0;

    return arena;
}

void arena_del(arena_t *arena)
{
    size_t msz = arena->a_sz + sizeof(arena_t);

    arena_pull(arena, SIZE_MAX);
    memset(arena, 0, sizeof(*arena));
    if (msz < sysconf(_SC_PAGE_SIZE))
    {
        free(arena);
    }
    else if (munmap(arena, msz) != 0)
    {
        fprintf(stderr, "munmap() failed: %s\n", strerror(errno));
    }
}

/*
 * Extend the arena boundary by `sz` bytes.
 * Can be used to grow the last allocated element.
 */
void *arena_push(arena_t *arena, size_t sz)
{
    void *ret = &arena->a_data[arena->a_pos];

    if ((SIZE_MAX - arena->a_pos) < sz ||
            arena->a_pos + sz > arena->a_sz)
    {
        return NULL;
    }

    arena->a_pos += sz;
    return ret;
}

/*
 * Reduce the arena boundary by `sz` bytes.
 */
void arena_pull(arena_t *arena, size_t sz)
{
    if (arena->a_pos < sz)
    {
        arena->a_pos = 0;
    }
    else
    {
        arena->a_pos -= sz;
    }
    arena_defer_cleanup(arena);
}

/*
 * Save current arena boundary, calling arena_restore() afterwards will
 * free everything that was allocated between arena_save() and arena_restore()
 */
arena_frame_t arena_save(arena_t *arena)
{
    if (arena == NULL) return (arena_frame_t){ .af_arena = NULL };
    return (arena_frame_t){ .af_arena = arena, .af_pos = arena->a_pos };
}

/*
 * Restore saved arena boundary (free elements between arena_save() and arena_restore())
 */
void arena_restore(arena_frame_t *frame)
{
    if (frame->af_arena == NULL) return;
    if (frame->af_pos > frame->af_arena->a_pos) return;
    arena_pull(frame->af_arena, frame->af_arena->a_pos - frame->af_pos);
}

/*
 * Allocate a defer action on the current arena. If the defer action block
 * is freed, the defer callback is called.
 */
bool arena_defer(arena_t *arena, arena_defer_fn_t *fn, void *data)
{
    struct arena_defer *defer;

    defer = arena_malloc(arena, sizeof(*defer));
    if (defer == NULL)
    {
        fprintf(stderr, "Unable to allocate defer buffer.");
        fn(arena, data);
        return false;
    }

    defer->ad_fn = fn;
    defer->ad_data = data;

    defer->ad_next = arena->a_defer;
    defer->ad_magic = ARENA_MAGIC;
    defer->ad_magic ^= (uintptr_t)defer;
    defer->ad_magic ^= (uintptr_t)defer->ad_next;
    defer->ad_magic ^= (uintptr_t)defer->ad_fn;
    defer->ad_magic ^= (uintptr_t)defer->ad_data;
    arena->a_defer = defer;

    return true;
}

void arena_defer_cleanup(arena_t *arena)
{
    while (arena->a_defer != NULL)
    {
        if (((uint8_t *)arena->a_defer + sizeof(struct arena_defer)) <= &arena->a_data[arena->a_pos])
            break;

        uintptr_t magic = arena->a_defer->ad_magic;
        magic ^= (uintptr_t)arena->a_defer;
        magic ^= (uintptr_t)arena->a_defer->ad_next;
        magic ^= (uintptr_t)arena->a_defer->ad_fn;
        magic ^= (uintptr_t)arena->a_defer->ad_data;
        if (magic != ARENA_MAGIC)
        {
            assert(!"Defer buffer corrupted.");
        }

        arena->a_defer->ad_fn(arena, arena->a_defer->ad_data);
        arena->a_defer = arena->a_defer->ad_next;
    }
}

void *arena_malign(arena_t *arena, size_t sz)
{
    size_t off;

    size_t align = 1;

    /* Calculate the alignment of sz */
    while (align < sz)
    {
        align <<= 1;
        if (align >= sizeof(max_align_t)) break;
    }

    /* Calculate the offset we need to add to the current position */
    off = align;
    off -= arena->a_pos & (align - 1);
    off &= (align - 1);
    if (arena_push(arena, off) == NULL) return NULL;

    return &arena->a_data[arena->a_pos];
}

void *arena_malloc(arena_t *arena, size_t sz)
{
    if (arena_malign(arena, sz) == NULL) return NULL;
    return arena_push(arena, sz);
}

/*
 * Allocate a string on the arena
 */
char *arena_strdup(arena_t *arena, const char *src)
{
    char *dst;

    size_t slen = strlen(src);

    dst = arena_push(arena, slen);
    if (dst == NULL) return NULL;

    return strcpy(dst, src);
}