#include <arena.h>
#include <liballoc.h>
#include <stddef.h>

static struct arena_chunk* arena_create_chunk (size_t capacity) {
  size_t size = sizeof (struct arena_chunk) + sizeof (uintptr_t) * capacity;
  struct arena_chunk* chunk = malloc (size);

  if (chunk == NULL)
    return NULL;

  chunk->next = NULL;
  chunk->size = 0;
  chunk->capacity = capacity;
  return chunk;
}

static void arena_destroy_chunk (struct arena_chunk* chunk) { free (chunk); }

void arena_reset (struct arena* arena) {
  for (struct arena_chunk* chunk = arena->begin; chunk != NULL; chunk = chunk->next) {
    chunk->size = 0;
  }

  arena->end = arena->begin;
}

void arena_destroy (struct arena* arena) {
  struct arena_chunk* chunk = arena->begin;

  while (chunk) {
    struct arena_chunk* chunk1 = chunk;
    chunk = chunk->next;
    arena_destroy_chunk (chunk1);
  }

  arena->begin = NULL;
  arena->end = NULL;
}

void* arena_malloc (struct arena* arena, size_t size) {
  size_t size1 = (size + sizeof (uintptr_t) - 1) / sizeof (uintptr_t);

  if (arena->end == NULL) {
    size_t capacity = ARENA_CHUNK_CAPACITY;

    if (capacity < size1) {
      capacity = size1;
    }

    arena->end = arena_create_chunk (capacity);

    if (arena->end == NULL) {
      return NULL;
    }

    arena->begin = arena->end;
  }

  while (arena->end->size + size1 > arena->end->capacity && arena->end->next != NULL) {
    arena->end = arena->end->next;
  }

  if (arena->end->size + size1 > arena->end->capacity) {
    size_t capacity = ARENA_CHUNK_CAPACITY;

    if (capacity < size1) {
      capacity = size1;
    }

    arena->end = arena_create_chunk (capacity);
    arena->end = arena->end->next;
  }

  void* result = &arena->end->memory[arena->end->size];
  arena->end->size += size1;
  return result;
}

void* arena_realloc (struct arena* arena, void* memory, size_t prev_size, size_t new_size) {
  if (new_size <= prev_size)
    return memory;

  void* new_memory = arena_malloc (arena, new_size);

  for (size_t i = 0; i < prev_size; i++)
    ((uint8_t*)new_memory)[i] = ((uint8_t*)memory)[i];

  return new_memory;
}