#include "../arena.h"
#include "../os.h"
#include "../core.h"

#include "runner.h"

UNIT_TEST(arena_init_virtual) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    ASSERT(arena.type == ARENA_VIRTUAL);
    ASSERT(arena.beg != NULL);
    ASSERT(arena.cur == arena.beg);
    ASSERT(arena.end == arena.beg + MB(1));
    arena_cleanup(&arena);
}

UNIT_TEST(arena_init_malloc) {
    arena_t arena = arena_make(ARENA_MALLOC, KB(4));
    ASSERT(arena.type == ARENA_MALLOC);
    ASSERT(arena.beg != NULL);
    ASSERT(arena.cur == arena.beg);
    ASSERT(arena.end == arena.beg + KB(4));
    arena_cleanup(&arena);
}

UNIT_TEST(arena_init_malloc_always) {
    arena_t arena = arena_make(ARENA_MALLOC_ALWAYS, KB(4));
    ASSERT(arena.type == ARENA_MALLOC_ALWAYS);
    arena_cleanup(&arena);
}

UNIT_TEST(arena_init_static) {
    u8 buffer[KB(4)];
    arena_t arena = arena_make(ARENA_STATIC, KB(4), buffer);
    ASSERT(arena.type == ARENA_STATIC);
    ASSERT(arena.beg == buffer);
    ASSERT(arena.cur == arena.beg);
    ASSERT(arena.end == arena.beg + KB(4));
    arena_cleanup(&arena);
}

UNIT_TEST(arena_alloc_basic) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    int *ptr = alloc(&arena, int);
    ASSERT(ptr != NULL);
    *ptr = 42;
    ASSERT(*ptr == 42);
    ASSERT(arena.cur > arena.beg);
    arena_cleanup(&arena);
}

UNIT_TEST(arena_alloc_array) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    int *arr = alloc(&arena, int, .count = 10);
    ASSERT(arr != NULL);
    for (int i = 0; i < 10; i++) {
        arr[i] = i;
    }
    for (int i = 0; i < 10; i++) {
        ASSERT(arr[i] == i);
    }
    arena_cleanup(&arena);
}

UNIT_TEST(arena_alloc_custom_align) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    void *ptr = alloc(&arena, char, .align = 64);
    ASSERT(ptr != NULL);
    ASSERT(((uintptr_t)ptr & 63) == 0); // Should be 64-byte aligned
    arena_cleanup(&arena);
}

UNIT_TEST(arena_alloc_nozero) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    int *ptr1 = alloc(&arena, int);
    ASSERT(*ptr1 == 0); // Default zeroed
    
    int *ptr2 = alloc(&arena, int, .flags = ALLOC_NOZERO);
    // We can't assert on the value as it's uninitialized
    ASSERT(ptr2 != NULL);
    arena_cleanup(&arena);
}

UNIT_TEST(arena_alloc_soft_fail) {
    u8 buffer[10];
    arena_t arena = arena_make(ARENA_STATIC, 10, buffer);
    
    void *ptr1 = alloc(&arena, char, .count = 5);
    ASSERT(ptr1 != NULL);
    
    // This would normally fail, but with SOFT_FAIL it returns NULL
    void *ptr2 = alloc(&arena, char, .count = 10, .flags = ALLOC_SOFT_FAIL);
    ASSERT(ptr2 == NULL);
    
    arena_cleanup(&arena);
}

UNIT_TEST(arena_scratch) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    arena_t scratch = arena_scratch(&arena, KB(1));
    
    ASSERT(scratch.beg != NULL);
    ASSERT(scratch.type == ARENA_STATIC);
    
    void *ptr = alloc(&scratch, int);
    ASSERT(ptr != NULL);
    
    // Scratch cleanup happens implicitly when parent arena is cleaned up
    arena_cleanup(&arena);
}

UNIT_TEST(arena_tell) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    usize pos1 = arena_tell(&arena);
    ASSERT(pos1 == 0);
    
    alloc(&arena, int);
    usize pos2 = arena_tell(&arena);
    ASSERT(pos2 > pos1);
    
    arena_cleanup(&arena);
}

UNIT_TEST(arena_remaining) {
    arena_t arena = arena_make(ARENA_VIRTUAL, KB(64));
    usize initial_remaining = arena_remaining(&arena);
    ASSERT(initial_remaining == KB(64));
    
    alloc(&arena, char, .count = KB(4));
    usize after_alloc = arena_remaining(&arena);
    ASSERT(after_alloc < initial_remaining);
    ASSERT(after_alloc >= KB(60)); // Account for possible alignment padding
    
    arena_cleanup(&arena);
}

UNIT_TEST(arena_capacity) {
    arena_t arena = arena_make(ARENA_VIRTUAL, KB(64));
    usize cap = arena_capacity(&arena);
    ASSERT(cap == KB(64));
    arena_cleanup(&arena);
}

UNIT_TEST(arena_rewind) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    
    usize mark = arena_tell(&arena);
    
    int *ptr1 = alloc(&arena, int);
    *ptr1 = 42;
    
    alloc(&arena, char, .count = 100);
    
    arena_rewind(&arena, mark);
    
    int *ptr2 = alloc(&arena, int);
    ASSERT(ptr2 == ptr1); // Should reuse the same memory
    
    // Original value is lost after rewind
    *ptr2 = 24;
    ASSERT(*ptr2 == 24);
    
    arena_cleanup(&arena);
}

UNIT_TEST(arena_pop) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    
    alloc(&arena, char, .count = 100);
    usize pos = arena_tell(&arena);
    
    alloc(&arena, char, .count = 50);
    
    arena_pop(&arena, 50);
    ASSERT(arena_tell(&arena) == pos);
    
    arena_cleanup(&arena);
}

UNIT_TEST(arena_malloc_arena) {
    void *ptr = alloc(&malloc_arena, int);
    ASSERT(ptr != NULL);
    
    // We need to free each allocation from malloc_arena manually
    os_free(ptr);
}

UNIT_TEST(arena_alloc_mixed_types) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(1));
    
    int *i = alloc(&arena, int);
    float *f = alloc(&arena, float);
    char *c = alloc(&arena, char);
    
    *i = 42;
    *f = 3.14f;
    *c = 'A';
    
    ASSERT(*i == 42);
    ASSERT(*f == 3.14f);
    ASSERT(*c == 'A');
    
    arena_cleanup(&arena);
}

UNIT_TEST(arena_multiple_arenas) {
    arena_t arena1 = arena_make(ARENA_VIRTUAL, KB(4));
    arena_t arena2 = arena_make(ARENA_VIRTUAL, KB(4));
    
    int *ptr1 = alloc(&arena1, int);
    int *ptr2 = alloc(&arena2, int);
    
    *ptr1 = 42;
    *ptr2 = 24;
    
    ASSERT(*ptr1 == 42);
    ASSERT(*ptr2 == 24);
    
    arena_cleanup(&arena1);
    arena_cleanup(&arena2);
}

UNIT_TEST(arena_stress_test) {
    arena_t arena = arena_make(ARENA_VIRTUAL, MB(10));
    
    // Allocate many objects
    for (int i = 0; i < 1000; i++) {
        int *ptr = alloc(&arena, int);
        ASSERT(ptr != NULL);
        *ptr = i;
    }
    
    // Allocate a large block
    void *large = alloc(&arena, char, .count = MB(5));
    ASSERT(large != NULL);
    
    arena_cleanup(&arena);
}