#define _GNU_SOURCE
#include "arena.h"

#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct block block;
struct block {
    block  *next;
    size_t  capacity;
    size_t  used;
};

struct arena {
    block  *head;
    block  *first;
    size_t  total_used;
    size_t  total_cap;
};

static block *block_create(size_t capacity) {
    block *b = malloc(sizeof(block) + capacity);
    if (b == nullptr) return nullptr;
    b->next = nullptr;
    b->capacity = capacity;
    b->used = 0;
    return b;
}

/**
 * arena_create() - Allocate a new arena with one initial block.
 * @initial_capacity: Bytes of usable space in the first block.
 *
 * Return: Arena pointer, or NULL on allocation failure.
 */
arena *arena_create(size_t initial_capacity) {
    arena *a = malloc(sizeof(arena));
    if (a == nullptr) return nullptr;

    block *b = block_create(initial_capacity);
    if (b == nullptr) {
        free(a);
        return nullptr;
    }

    a->head = b;
    a->first = b;
    a->total_used = 0;
    a->total_cap = initial_capacity;
    return a;
}

/**
 * arena_destroy() - Free every block held by the arena.
 * @a: Arena to free. NULL is permitted.
 */
void arena_destroy(arena *a) {
    if (a == nullptr) return;
    block *b = a->first;
    while (b != nullptr) {
        block *next = b->next;
        free(b);
        b = next;
    }
    free(a);
}

/**
 * arena_alloc() - Bump-allocate an aligned chunk from the arena.
 * @a: Arena to allocate from.
 * @size: Number of bytes to reserve.
 * @align: Required alignment in bytes (power of two).
 *
 * A new block at twice the current capacity (or larger if the request
 * itself is larger) is allocated when the head block is exhausted.
 *
 * Return: Pointer to @size aligned bytes, or NULL on allocation failure.
 */
void *arena_alloc(arena *a, size_t size, size_t align) {
    if (a == nullptr || size == 0) return nullptr;

    uintptr_t base = (uintptr_t)((char *)(a->head + 1) + a->head->used);
    uintptr_t aligned = (base + (align - 1)) & ~(align - 1);
    size_t pad = aligned - base;

    if (a->head->used + pad + size <= a->head->capacity) {
        a->head->used += pad + size;
        a->total_used += pad + size;
        return (void *)aligned;
    }

    size_t new_cap = a->head->capacity * 2;
    if (new_cap < size + align) new_cap = size + align;

    block *b = block_create(new_cap);
    if (b == nullptr) return nullptr;

    b->next = a->head;
    a->head = b;
    a->total_cap += new_cap;

    base = (uintptr_t)(b + 1);
    aligned = (base + (align - 1)) & ~(align - 1);
    pad = aligned - base;

    b->used = pad + size;
    a->total_used += pad + size;
    return (void *)aligned;
}

/**
 * arena_strdup() - Duplicate a NUL-terminated string into the arena.
 * @a: Arena to allocate from.
 * @s: Source string. NULL yields NULL.
 *
 * Return: Owned-by-arena copy of @s, or NULL on allocation failure.
 */
char *arena_strdup(arena *a, const char *s) {
    if (s == nullptr) return nullptr;
    size_t len = strlen(s) + 1;
    char *out = arena_alloc(a, len, 1);
    if (out == nullptr) return nullptr;
    memcpy(out, s, len);
    return out;
}

/**
 * arena_sprintf() - printf into a fresh arena allocation.
 * @a: Arena to allocate from.
 * @fmt: printf format string.
 *
 * Return: Owned-by-arena formatted string, or NULL on allocation failure.
 */
char *arena_sprintf(arena *a, const char *fmt, ...) {
    va_list ap;

    va_start(ap, fmt);
    int len = vsnprintf(nullptr, 0, fmt, ap);
    va_end(ap);
    if (len < 0) return nullptr;

    char *out = arena_alloc(a, (size_t)len + 1, 1);
    if (out == nullptr) return nullptr;

    va_start(ap, fmt);
    vsnprintf(out, (size_t)len + 1, fmt, ap);
    va_end(ap);
    return out;
}

/**
 * arena_reset() - Rewind the arena to empty without releasing memory.
 * @a: Arena to reset. NULL is permitted.
 *
 * Frees every block except the first; the first block's used pointer is
 * reset to zero.
 */
void arena_reset(arena *a) {
    if (a == nullptr) return;

    block *b = a->head;
    while (b != a->first) {
        block *next = b->next;
        free(b);
        b = next;
    }
    a->head = a->first;
    a->first->used = 0;
    a->total_used = 0;
    a->total_cap = a->first->capacity;
}

size_t arena_used(const arena *a)     { return a == nullptr ? 0 : a->total_used; }
size_t arena_capacity(const arena *a) { return a == nullptr ? 0 : a->total_cap; }