#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 <util/util.h>
#include <assert.h>

#define ARENA_SIZE (1<<20)
#define BLOCK_MAGIC 0xDEADBEEF

typedef struct UmArena {
  struct UmArena *next;
  uintptr_t mem;
  uintptr_t base;
  size_t cursor;
  int64_t balance;
} __attribute__((aligned(16))) UmArena;

typedef struct {
  uint32_t magic;
  size_t size;
  uint8_t data[];
} __attribute__((aligned(16))) UmBlock;

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;
  }

  UmArena *arena = (UmArena *)mem;
  arena->base = (uintptr_t)mem;
  arena->mem = ((uintptr_t)mem + sizeof(UmArena) + 15) & ~0xF;
  arena->cursor = 0;
  arena->balance = 0;
  LL_APPEND(ARENAS, arena);

  return arena;
}

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

  size = (size + 15) & ~(size_t)0xF;

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

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

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

  uintptr_t addr = usable->mem + usable->cursor;
  UmBlock *block = (UmBlock *)addr;
  block->size = size;
  block->magic = BLOCK_MAGIC;

  usable->cursor += sizeof(UmBlock) + size;
  usable->balance += size;

  string_memset(block->data, 0, size);
  return block->data;
}

void ufree(void *ptr_) {
  if (ptr_ == NULL) {
    return;
  }

  uintptr_t ptr = (uintptr_t)ptr_;

  UmArena *freeable = NULL;
  size_t size = 0;

  UmArena *arena, *arenatmp;
  LL_FOREACH_SAFE(ARENAS, arena, arenatmp) {
    if (ptr >= arena->mem && ptr < arena->mem + arena->cursor) {
      UmBlock *block = (UmBlock *)(ptr - sizeof(UmBlock));

      if (((uintptr_t)block->data != ptr)
       || (block->magic != BLOCK_MAGIC)
       || (block->size == 0)
       || (block->size > ARENA_SIZE)) {
        return;
      }

      string_memset(block->data, 0xDD, block->size);
      block->magic = 0;

      arena->balance -= block->size;
      ASSERT(arena->balance >= 0, "umalloc: imbalance after free\n");

      if (arena->balance == 0) {
        LL_REMOVE(ARENAS, arena);
        mman_unmap((void *)arena->base);
      }

      return;
    }
  }
}