#include "src/memory.h"
#include "src/panic.h"
#include "src/util.h"
#include <stdint.h>

#define MEMORY_SIZE 134217728
#define PAGE_SIZE 4096
#define PAGE_COUNT (MEMORY_SIZE / PAGE_SIZE)
#define MAP_UNIT 64
#define FULL_UNIT 0xFFFFFFFFFFFFFFFF

// Claim 2MB for bootloader, kernel, and stuff.
static uint16_t free_page = MAP_UNIT * 8;
static uint64_t allocation[PAGE_COUNT / MAP_UNIT] = {FULL_UNIT, FULL_UNIT, FULL_UNIT, FULL_UNIT,
                                                     FULL_UNIT, FULL_UNIT, FULL_UNIT, FULL_UNIT};

static uint8_t get_allocated(uint16_t page) {
  return !!allocation[page / MAP_UNIT] & ((uint64_t)1 << (page % MAP_UNIT));
}

static void set_allocated(uint16_t page, uint8_t allocated) {
  if (allocated) {
    allocation[page / MAP_UNIT] |= ((uint64_t)1 << (page % MAP_UNIT));
  } else {
    allocation[page / MAP_UNIT] &= ~((uint64_t)1 << (page % MAP_UNIT));
  }
}

static void *memory_page_allocate() {
  ASSERT(free_page < PAGE_COUNT, "memory_page_allocate: out of memory");
  const uint16_t page = free_page;
  set_allocated(page, 1);
  for (uint16_t unit = free_page / MAP_UNIT; unit < PAGE_COUNT / MAP_UNIT; unit++) {
    if (allocation[unit] != FULL_UNIT) {
      uint16_t bit = (uint16_t)__builtin_ctzll(~allocation[unit]);
      free_page = unit * MAP_UNIT + bit;
      break;
    }
  }
  return (void *)((uint64_t)page * PAGE_SIZE);
}

__attribute__((unused)) static void memory_page_free(void *address) {
  uint16_t page = (uint16_t)((uint64_t)address / PAGE_SIZE);
  ASSERT(get_allocated(page), "memory_page_free: page not allocated")
  set_allocated(page, 0);
  if (page < free_page) {
    free_page = page;
  }
}

#define HEAP_PAGE_COUNT 256

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

static free_chunk_t *free_list;

void memory_init() {
  free_list = memory_page_allocate();
  for (uint64_t i = 1; i < HEAP_PAGE_COUNT; i++) {
    memory_page_allocate(); // TODO Map the pages.
  }
  free_list->size = HEAP_PAGE_COUNT * PAGE_SIZE;
  free_list->next = NUL;
}

void *memory_allocate(uint64_t size) {
  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;
  }

  ASSERT(0, "memory_heap_allocate: out of heap memory");
  return 0;
}

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