#ifndef ARENA_H
#define ARENA_H

#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

typedef struct ArenaBlock {
  struct ArenaBlock* next;
  size_t used;
  size_t capacity;
  max_align_t data[];
} ArenaBlock;

typedef struct {
  ArenaBlock* first;
  ArenaBlock* current;
  size_t block_size;
} Arena;

static inline Arena arena_create(void) {
  return (Arena){.block_size = 4096};
}

static inline void* arena_alloc_aligned(Arena* arena, size_t size,
                                        size_t alignment) {
  if (alignment == 0 || (alignment & (alignment - 1)) != 0) {
    return NULL;
  }
  if (size == 0) {
    size = 1;
  }

  ArenaBlock* block = arena->current;
  size_t offset = 0;
  while (block != NULL) {
    uintptr_t address = (uintptr_t)block->data + block->used;
    size_t padding = (size_t)(-(uintptr_t)address & (alignment - 1));
    if (padding <= SIZE_MAX - block->used) {
      offset = block->used + padding;
      if (offset <= block->capacity && size <= block->capacity - offset) {
        break;
      }
    }
    block = block->next;
  }

  if (block == NULL) {
    size_t capacity = arena->block_size == 0 ? 4096 : arena->block_size;
    if (size > SIZE_MAX - alignment) {
      return NULL;
    }
    if (capacity < size + alignment) {
      capacity = size + alignment;
    }
    if (capacity > SIZE_MAX - sizeof(ArenaBlock)) {
      return NULL;
    }

    ArenaBlock* new_block = malloc(sizeof(*new_block) + capacity);
    if (new_block == NULL) {
      return NULL;
    }
    *new_block = (ArenaBlock){.capacity = capacity};
    ArenaBlock* previous = arena->current;
    if (previous == NULL) {
      arena->first = new_block;
    } else {
      while (previous->next != NULL) {
        previous = previous->next;
      }
      previous->next = new_block;
    }
    arena->current = new_block;
    block = new_block;
    uintptr_t address = (uintptr_t)block->data;
    offset = (size_t)(-(uintptr_t)address & (alignment - 1));
  } else {
    arena->current = block;
  }

  void* result = (unsigned char* )block->data + offset;
  block->used = offset + size;
  return result;
}

static inline void* arena_alloc(Arena* arena, size_t size) {
  return arena_alloc_aligned(arena, size, _Alignof(max_align_t));
}

static inline void arena_reset(Arena* arena) {
  for (ArenaBlock* block = arena->first; block != NULL; block = block->next) {
    block->used = 0;
  }
  arena->current = arena->first;
}

static inline void arena_destroy(Arena* arena) {
  ArenaBlock* block = arena->first;
  while (block != NULL) {
    ArenaBlock* next = block->next;
    free(block);
    block = next;
  }
  *arena = (Arena){0};
}

#endif