#include <stdint.h>
#include <stddef.h>
#include <umalloc/umalloc.h>
#include <sysdefs/mman.h>
#include <system/system.h>
#include <errors.h>
#include <linklist.h>
#include <string/string.h>
#include <write/write.h>
#include <util/util.h>

#define ARENA_SIZE (1<<20)

typedef struct UmArena {
  struct UmArena *next;
  uintptr_t mem;
  size_t cursor;
  int balance;
} UmArena;

static UmArena *ARENAS = NULL;

UmArena *um_newarena(void) {
  size_t arenasize = ARENA_SIZE;

  uint8_t *mem = NULL;
  int32_t err = mman_map(NULL, arenasize, MMAN_MAP_PF_RW, 0, &mem);
  if (mem == NULL || err != E_OK) {
    return NULL;
  }
  string_memset(mem, 0, arenasize);

  UmArena *arena = (UmArena *)mem;
  arena->mem = (uintptr_t)mem + sizeof(UmArena);
  arena->cursor = 0;
  LL_APPEND(ARENAS, arena);

  return arena;
}

void *umalloc(size_t size) {
  UmArena *usable = NULL;

  UmArena *arena, *arenatmp;
  LL_FOREACH_SAFE(ARENAS, arena, arenatmp) {
    if (arena->cursor + size <= ARENA_SIZE - sizeof(UmArena)) {
      usable = arena;
      break;
    }
  }

  if (usable == NULL) {
    usable = um_newarena();
  }

  if (usable == NULL) {
    return NULL;
  }

  uintptr_t current = usable->mem + usable->cursor;
  usable->cursor += size;
  usable->balance += 1;
  return (void *)current;
}

void ufree(void *ptr_) {
  uintptr_t ptr = (uintptr_t)ptr_;

  UmArena *freeable = NULL;

  UmArena *arena, *arenatmp;
  LL_FOREACH_SAFE(ARENAS, arena, arenatmp) {
    if (ptr >= arena->mem && ptr < arena->mem + ARENA_SIZE - sizeof(UmArena)) {
      freeable = arena;
      break;
    }
  }

  if (freeable == NULL) {
    return;
  }

  freeable->balance -= 1;

  if (freeable->balance < 0) {
    writefmt("umalloc: ufree() called more times than umalloc() ({d})\n", freeable->balance);
    quit();
  }

  if (freeable->balance == 0) {
    LL_REMOVE(ARENAS, freeable);

    mman_unmap((void *)(freeable->mem - sizeof(UmArena)));
  }
}