#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#define ARENA_IMPLEMENTATION
#define ARENA_NO_ASSERT
#include "arena.h"

// Helper to count active regions in the linked list
static int count_regions(const Arena *arena)
{
    int count = 0;
    for (Arena__Region *r = arena->head; r != NULL; r = r->next) {
        count++;
    }
    return count;
}

// Helper to verify pointer alignment
static bool is_aligned(const void *ptr)
{
    return ((uintptr_t)ptr % ARENA_ALIGNMENT) == 0;
}

static void test_edge_case_inputs(void)
{
    printf("Running edge case and invalid input tests...\n");
    Arena arena = arena_create(1024);

    // 0-byte allocations
    assert(arena_alloc(&arena, 0) == NULL);
    assert(arena_calloc(&arena, 0, 10) == NULL);
    assert(arena_calloc(&arena, 10, 0) == NULL);

    // NULL pointer handling
    assert(arena_alloc(NULL, 100) == NULL);
    assert(arena_calloc(NULL, 1, 1) == NULL);
    assert(arena_copy(NULL, "data", 4) == NULL);
    assert(arena_copy(&arena, NULL, 4) == NULL);
    assert(arena_strdup(NULL, "test") == NULL);
    assert(arena_strdup(&arena, NULL) == NULL);
    assert(arena_sprintf(NULL, "test") == NULL);
    assert(arena_sprintf(&arena, NULL) == NULL);

    // Snapshot/Rewind/Trim edge cases on NULL or empty
    arena_trim(NULL); // Should not crash
    arena_reset(NULL); // Should not crash
    arena_free(NULL); // Should not crash

    ArenaMarker null_marker = arena_snapshot(NULL);
    assert(null_marker.region == NULL && null_marker.count == 0);
    arena_rewind(NULL, null_marker); // Should not crash

    arena_free(&arena);
}

static void test_overflows(void)
{
    printf("Running integer overflow tests...\n");
    Arena arena = arena_create(1024);

    // calloc multiplication overflow (count * size > SIZE_MAX)
    size_t massive_count = SIZE_MAX / 2;
    size_t massive_size = 4;
    assert(arena_calloc(&arena, massive_count, massive_size) == NULL);

    // arena_alloc size wrap-around on alignment
    // arena__align(SIZE_MAX) will overflow back to a small number
    size_t overflow_align_size = SIZE_MAX - (ARENA_ALIGNMENT / 2);
    assert(arena_alloc(&arena, overflow_align_size) == NULL);

    arena_free(&arena);
}

static void test_geometric_growth_and_stress(void)
{
    printf("Running geometric growth and stress tests...\n");

    size_t initial_capacity = 64;
    Arena arena = arena_create(initial_capacity);

    // Exhaust the first region to force geometric growth
    assert(arena_alloc(&arena, 32) != NULL); // Head allocated with capacity 64
    assert(arena.head->capacity == 64);

    assert(arena_alloc(&arena, 32) != NULL); // Head is now full

    // Trigger next region (should be 64 * 2 = 128)
    assert(arena_alloc(&arena, 32) != NULL);
    assert(arena.tail->capacity == 128);

    // Trigger next region (should be 128 * 2 = 256)
    assert(arena_alloc(&arena, 128) != NULL);
    assert(arena.tail->capacity == 256);

    // Stress test: Allocate 10,000 small blocks
    for (int i = 0; i < 10000; i++) {
        void *ptr = arena_alloc(&arena, 7); // Odd size to test alignment scaling
        assert(ptr != NULL);
        assert(is_aligned(ptr));
    }

    // Verify regions were created
    int regions = count_regions(&arena);
    assert(regions > 3);

    // Huge allocation bypassing standard growth
    // Requesting 1MB when the next capacity might only be ~4KB
    void *huge = arena_alloc(&arena, 1024 * 1024);
    assert(huge != NULL);
    assert(arena.tail->capacity >= 1024 * 1024);

    arena_free(&arena);
}

static void test_snapshot_rewind_trim_logic(void)
{
    printf("Running snapshot, rewind, and trim logic tests...\n");
    Arena arena = arena_create(128);

    // 1. Snapshot empty
    ArenaMarker empty_marker = arena_snapshot(&arena);
    assert(empty_marker.region == NULL);

    // 2. Allocate across 3 regions
    arena_alloc(&arena, 100); // Region 1 (Cap: 128)
    arena_snapshot(&arena);

    arena_alloc(&arena, 200); // Region 2 (Cap: 256)
    ArenaMarker mark_r2 = arena_snapshot(&arena);

    arena_alloc(&arena, 400); // Region 3 (Cap: 512)
    assert(count_regions(&arena) == 3);

    // 3. Rewind to Region 2
    arena_rewind(&arena, mark_r2);
    assert(arena.tail == mark_r2.region);
    assert(arena.tail->count == mark_r2.count);

    // Check that Region 3's count was reset to 0
    assert(arena.tail->next != NULL);
    assert(arena.tail->next->count == 0);

    // 4. Trim trailing regions (Should drop Region 3)
    arena_trim(&arena);
    assert(arena.tail->next == NULL);
    assert(count_regions(&arena) == 2);

    // 5. Rewind to completely empty
    arena_rewind(&arena, empty_marker);
    assert(arena.tail == NULL || (arena.tail == arena.head && arena.tail->count == 0));

    // Trim when virtually empty (should retain the head region if it exists)
    arena_trim(&arena);
    if (arena.head != NULL) {
        assert(arena.head->next == NULL);
    }

    arena_free(&arena);
}

static void test_massive_string_formatting(void)
{
    printf("Running massive string formatting tests...\n");
    Arena arena = arena_create(1024);

    // Create a 50,000 character string requirement
    size_t massive_len = 50000;
    char *massive_input = malloc(massive_len + 1);
    assert(massive_input != NULL);
    memset(massive_input, 'X', massive_len);
    massive_input[massive_len] = '\0';

    char *result = arena_sprintf(&arena, "PREFIX_%s_SUFFIX", massive_input);
    assert(result != NULL);

    size_t expected_length = 7 + massive_len + 7;
    assert(strlen(result) == expected_length);
    assert(strncmp(result, "PREFIX_XXX", 10) == 0);
    assert(strcmp(result + expected_length - 7, "_SUFFIX") == 0);

    free(massive_input);
    arena_free(&arena);
}

static void test_zero_capacity_fallback(void)
{
    printf("Running zero-capacity fallback tests...\n");

    // Passing 0 should fall back to ARENA_REGION_CAPACITY implicitly on first alloc
    Arena arena = arena_create(0);

    void *ptr = arena_alloc(&arena, 10);
    assert(ptr != NULL);
    assert(arena.head->capacity >= ARENA_REGION_CAPACITY);

    arena_free(&arena);
}

int main(void)
{
    test_edge_case_inputs();
    test_overflows();
    test_geometric_growth_and_stress();
    test_snapshot_rewind_trim_logic();
    test_massive_string_formatting();
    test_zero_capacity_fallback();

    return EXIT_SUCCESS;
}