#include <cmocka.h>
#include <limits.h>
#include <setjmp.h>
#include <stdalign.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>

#include "lae_allocator.h"
#include "lae_arena.h"

static void test_push_3_ints(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(
        arena_init(&arena, allocator_default(), sizeof(int) * 3),
        ARENA_OK);

    int int_to_push = 20;
    ArenaPointer result =
        arena_push_data(arena, &int_to_push, sizeof(int), alignof(int));
    assert_true(result.is_valid);
    assert_int_equal(20, *(int *)arena_unwrap_pointer(result));
    assert_int_equal(20, *(int *)arena_unwrap_pointer(result));

    int_to_push = 30;
    result = arena_push_data(arena, &int_to_push, sizeof(int), alignof(int));
    assert_true(result.is_valid);
    assert_int_equal(30, *(int *)arena_unwrap_pointer(result));

    int_to_push = 40;
    result = arena_push_data(arena, &int_to_push, sizeof(int), alignof(int));
    assert_true(result.is_valid);
    assert_int_equal(40, *(int *)arena_unwrap_pointer(result));

    arena_destroy(&arena);
}

static void test_push_3_ints_2_doubles(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(
        arena_init(
            &arena,
            allocator_default(),
            (sizeof(int) * 3) + (sizeof(double) * 2)),
        ARENA_OK);

    int int_to_push = 20;
    ArenaPointer result =
        arena_push_data(arena, &int_to_push, sizeof(int), alignof(int));
    assert_true(result.is_valid);
    assert_int_equal(20, *(int *)arena_unwrap_pointer(result));

    double double_to_push = 4.57;
    result = arena_push_data(
        arena,
        &double_to_push,
        sizeof(double),
        alignof(double));
    assert_true(result.is_valid);
    assert_double_equal(4.57, *(double *)arena_unwrap_pointer(result), 1e-6);

    int_to_push = 30;
    result = arena_push_data(arena, &int_to_push, sizeof(int), alignof(int));
    assert_true(result.is_valid);
    assert_int_equal(30, *(int *)arena_unwrap_pointer(result));

    int_to_push = 40;
    result = arena_push_data(arena, &int_to_push, sizeof(int), alignof(int));
    assert_true(result.is_valid);
    assert_int_equal(40, *(int *)arena_unwrap_pointer(result));

    double_to_push = 267.33;
    result = arena_push_data(
        arena,
        &double_to_push,
        sizeof(double),
        alignof(double));
    assert_true(result.is_valid);
    assert_double_equal(267.33, *(double *)arena_unwrap_pointer(result), 1e-6);

    arena_destroy(&arena);
}

static void test_init_arena_0_cap(void **state) {
    (void)state;

    Arena *arena;
    assert_int_equal(
        arena_init(&arena, allocator_default(), 0),
        ARENA_INVALID_SIZE);
}

static void test_arena_alloc_size_0(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaPointer pointer = arena_alloc(arena, 0, alignof(int));
    assert_false(pointer.is_valid);
    assert_int_equal(pointer.err, ARENA_INVALID_SIZE);

    arena_destroy(&arena);
}

static void test_arena_alloc_size_max(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaPointer pointer = arena_alloc(arena, SIZE_MAX, alignof(int));
    assert_false(pointer.is_valid);
    assert_int_equal(pointer.err, ARENA_OUT_OF_SPACE);

    arena_destroy(&arena);
}

static void test_arena_alloc_align_0(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaPointer pointer = arena_alloc(arena, sizeof(int), 0);
    assert_false(pointer.is_valid);
    assert_int_equal(pointer.err, ARENA_INVALID_ALIGN);

    arena_destroy(&arena);
}

static void test_arena_alloc_align_max(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaPointer pointer = arena_alloc(arena, sizeof(int), SIZE_MAX);
    assert_false(pointer.is_valid);
    assert_int_equal(pointer.err, ARENA_INVALID_ALIGN);

    arena_destroy(&arena);
}

static void test_arena_align_size_0(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaPointer pointer = arena_alloc(arena, sizeof(int), 0);
    assert_false(pointer.is_valid);
    assert_int_equal(pointer.err, ARENA_INVALID_ALIGN);

    arena_destroy(&arena);
}

static void test_arena_ensure_capacity_size_0(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaErr err = arena_ensure_capacity(arena, 0, alignof(int));
    assert_int_equal(err, ARENA_INVALID_SIZE);

    arena_destroy(&arena);
}

static void test_arena_ensure_capacity_size_max(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaErr err = arena_ensure_capacity(arena, SIZE_MAX, alignof(int));
    assert_int_equal(err, ARENA_INVALID_SIZE);

    arena_destroy(&arena);
}

static void test_arena_ensure_capacity_align_0(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaErr err = arena_ensure_capacity(arena, sizeof(int), 0);
    assert_int_equal(err, ARENA_INVALID_ALIGN);

    arena_destroy(&arena);
}

static void test_arena_ensure_capacity_align_max(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaErr err = arena_ensure_capacity(arena, sizeof(int), SIZE_MAX);
    assert_int_equal(err, ARENA_INVALID_ALIGN);

    arena_destroy(&arena);
}

static void test_arena_allocator_valid_allocation(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaAllocator result = arena_get_allocator(arena);
    assert_true(result.is_valid);

    Allocator arena_allocator = result.allocator;

    int *x =
        arena_allocator.alloc(arena_allocator.ctx, sizeof(int), alignof(int));

    assert_non_null(x);

    *x = 34;
    assert_int_equal(*x, 34);

    arena_destroy(&arena);
}

static void test_arena_allocatro_many_allocations(void **state) {
    (void)state;

    Arena *arena;
    assert_uint_equal(arena_init(&arena, allocator_default(), 64), ARENA_OK);

    ArenaAllocator result = arena_get_allocator(arena);
    assert_true(result.is_valid);

    Allocator arena_allocator = result.allocator;
    for (int i = 0; i < 10; i++) {
        int *x = arena_allocator.alloc(
            arena_allocator.ctx,
            sizeof(int),
            alignof(int));

        assert_non_null(x);
    }

    arena_destroy(&arena);
}

static void test_arena_allocator_from_null_arena(void **state) {
    (void)state;

    Arena *arena = NULL;

    ArenaAllocator result = arena_get_allocator(arena);
    assert_false(result.is_valid);
}

int main(void) {
    const struct CMUnitTest tests[] = {
        cmocka_unit_test(test_push_3_ints),
        cmocka_unit_test(test_push_3_ints_2_doubles),
        cmocka_unit_test(test_init_arena_0_cap),
        cmocka_unit_test(test_arena_alloc_size_0),
        cmocka_unit_test(test_arena_alloc_size_max),
        cmocka_unit_test(test_arena_alloc_align_0),
        cmocka_unit_test(test_arena_alloc_align_max),
        cmocka_unit_test(test_arena_align_size_0),
        cmocka_unit_test(test_arena_ensure_capacity_align_0),
        cmocka_unit_test(test_arena_ensure_capacity_align_max),
        cmocka_unit_test(test_arena_ensure_capacity_size_0),
        cmocka_unit_test(test_arena_ensure_capacity_size_max),
        cmocka_unit_test(test_arena_allocator_valid_allocation),
    };

    return cmocka_run_group_tests(tests, NULL, NULL);
}