#include "arena.h"
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

ArenaResult arena_init(size_t capacity) {

    void *buffer = malloc(capacity);
    if (buffer == NULL) {
        ArenaResult err = {
            .is_valid = false,
            .err = ARENA_BAD_ALLOC,
        };
        return err;
    }

    if (capacity < 1) {
        ArenaResult err = {
            .is_valid = false,
            .err = ARENA_INVALID_SIZE,
        };
        return err;
    }

    Arena new_arena = {
        .buffer = buffer,
        .capacity = capacity,
        .offset = 0,
    };
    ArenaResult val = {
        .is_valid = true,
        .arena = new_arena,
    };

    return val;
}

void arena_destroy(Arena *arena) {
    if (arena == NULL) {
        return;
    }

    arena->offset = 0;
    arena->capacity = 0;
    free(arena->buffer);
}

ArenaPointer arena_alloc(Arena *arena, size_t size, size_t alignment) {
    if (arena == NULL) {
        ArenaPointer err = {
            .is_valid = false,
            .err = ARENA_NULL_ARG,
        };
        return err;
    }

    SizeResult new_offset = align_arena_offset(arena, alignment);
    if (!new_offset.is_valid) {
        ArenaPointer err = {
            .is_valid = false,
            .err = new_offset.err,
        };
        return err;
    }

    if (new_offset.val + size >= arena->capacity) {
        ArenaPointer err = {
            .is_valid = false,
            .err = ARENA_OUT_OF_SPACE,
        };
        return err;
    }

    arena->offset = new_offset.val;

    ArenaPointer val = {
        .is_valid = true,
        .address = arena->buffer + arena->offset,
    };
    return val;
}

ArenaErr arena_push(Arena *arena, void *data, size_t size, size_t alignment) {
    if (arena == NULL || data == NULL) {
        return ARENA_NULL_ARG;
    }

    ArenaPointer pointer = arena_alloc(arena, size, alignment);

    if (!pointer.is_valid) {
        return pointer.err;
    }
    while (!pointer.is_valid && pointer.err == ARENA_OUT_OF_SPACE) {
        size_t new_capacity;

        bool is_overflow = mul_size_t_safe(arena->capacity, 2, &new_capacity);
        if (is_overflow) {
            return ARENA_CAPACITY_OVERFLOW;
        }

        ArenaErr err = arena_realloc(arena, new_capacity);
        if (err != ARENA_OK) {
            return err;
        }

        pointer = arena_alloc(arena, size, alignment);
    }

    memcpy(
        pointer.address,
        data,
        size);

    return ARENA_OK;
}

ArenaErr arena_realloc(Arena *arena, size_t new_capacity) {
    if (arena == NULL) {
        return ARENA_NULL_ARG;
    }

    uint8_t *tmp = realloc(arena->buffer, new_capacity);
    if (tmp == NULL) {
        return ARENA_BAD_ALLOC;
    }

    arena->buffer = tmp;
    arena->capacity = new_capacity;
    return ARENA_OK;
}

SizeResult align_arena_offset(Arena *arena, size_t alignment) {
    if (arena == NULL) {
        SizeResult err = {
            .is_valid = false,
            .err = ARENA_NULL_ARG,
        };
        return err;
    }

    if (alignment < 1) {
        SizeResult err = {
            .is_valid = false,
            .err = ARENA_INVALID_ALIGN,
        };
        return err;
    }

    SizeResult val = {
        .is_valid = true,
        .val = ((arena->offset + (alignment - 1)) & ~(alignment - 1)),
        // Formula for alignment, found at stackoverflow, i came up with one
        // but it involved 3 divisions, this one is way faster :)
    };
    return val;
}

bool mul_size_t_safe(size_t a, size_t b, size_t *out) {
    if (out == NULL) {
        return true;
    }

    if (a != 0 && b > SIZE_MAX / a) {
        return true;
    }

    *out = a * b;
    return false;
}