#include <arena.h>
#include <malloc.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;
}