#include "src/lib/memory.h"
#include "src/lib/layout.h"
#include "src/lib/util.h"

#ifdef KERNEL
#define HEAP_VIRTUAL_BASE KERNEL_VIRTUAL_HEAP
#else
#define HEAP_VIRTUAL_BASE USER_VIRTUAL_HEAP
#endif

typedef struct free_chunk {
  uint64_t size; // header + data
  struct free_chunk *next;
} free_chunk_t;

static free_chunk_t *free_list;

void *memory_allocate(uint64_t size) {
  if (!free_list) {
    free_list = (free_chunk_t *)HEAP_VIRTUAL_BASE;
    free_list->size = HEAP_SIZE;
    free_list->next = NUL;
  }

  size = sizeof(free_chunk_t) + ((size + 7) & (uint64_t)~7);

  free_chunk_t *prev = NUL;
  free_chunk_t *curr = free_list;

  while (curr) {
    if (curr->size >= size + (sizeof(free_chunk_t) + 8)) {
      free_chunk_t *remainder = (free_chunk_t *)((uint8_t *)curr + size);
      remainder->size = curr->size - size;
      remainder->next = curr->next;
      curr->size = size;
      curr->next = remainder;
    }
    if (curr->size >= size) {
      if (prev) {
        prev->next = curr->next;
      } else {
        free_list = curr->next;
      }
      void *result = (void *)((uint8_t *)curr + sizeof(free_chunk_t));
      memory_set(0, size - sizeof(free_chunk_t), result);
      return result;
    }
    prev = curr;
    curr = curr->next;
  }

  return NUL;
}

void memory_free(void *pointer) {
  free_chunk_t *chunk = (free_chunk_t *)((uint8_t *)pointer - sizeof(free_chunk_t));
  chunk->next = free_list;
  free_list = chunk;
  // TODO Merge adjacent free chunks.
}