#include "src/kernel/memory.h"
#include "src/kernel/panic.h"
#include "src/lib/layout.h"
#include "src/lib/memory.h"

#define MAP_UNIT 64
#define FULL_UNIT 0xFFFFFFFFFFFFFFFF

// TODO keep aligned with `src/lib/layout.h`.
static uint16_t free_page = MAP_UNIT * 12;
static uint64_t allocation[PAGE_COUNT / MAP_UNIT] = {
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    FULL_UNIT,
    [59] = 0x8000000000000000ULL,
};

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

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

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

void memory_page_map(uint64_t *pml4, void *virt, void *phys, uint64_t flags) {
  const uint64_t ivirt = (uint64_t)virt;
  const uint64_t pml4_index = (ivirt >> 39) & 0x1FF;
  const uint64_t pdpt_index = (ivirt >> 30) & 0x1FF;
  const uint64_t pd_index = (ivirt >> 21) & 0x1FF;
  const uint64_t pt_index = (ivirt >> 12) & 0x1FF;

  if (!(pml4[pml4_index] & PAGE_PRESENT)) {
    uint64_t *pdpt = memory_page_allocate();
    memory_set(0, PAGE_SIZE, PHYS_TO_VIRT(pdpt));
    pml4[pml4_index] = (uint64_t)pdpt | PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER;
  }

  uint64_t *pdpt = PHYS_TO_VIRT(pml4[pml4_index] & ~(uint64_t)0xFFF);
  if (!(pdpt[pdpt_index] & PAGE_PRESENT)) {
    uint64_t *pd = memory_page_allocate();
    memory_set(0, PAGE_SIZE, PHYS_TO_VIRT(pd));
    pdpt[pdpt_index] = (uint64_t)pd | PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER;
  }

  uint64_t *pd = PHYS_TO_VIRT(pdpt[pdpt_index] & ~(uint64_t)0xFFF);
  ASSERT(!(pd[pd_index] & 0x80), "memory_page_map: huge page in PD");
  if (!(pd[pd_index] & PAGE_PRESENT)) {
    uint64_t *pt = memory_page_allocate();
    memory_set(0, PAGE_SIZE, PHYS_TO_VIRT(pt));
    pd[pd_index] = (uint64_t)pt | PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER;
  }

  uint64_t *pt = PHYS_TO_VIRT(pd[pd_index] & ~(uint64_t)0xFFF);
  pt[pt_index] = (uint64_t)phys | flags | PAGE_PRESENT;
}

void memory_page_unmap(uint64_t *pml4, void *virt) {
  const uint64_t ivirt = (uint64_t)virt;
  const uint64_t pml4_index = (ivirt >> 39) & 0x1FF;
  const uint64_t pdpt_index = (ivirt >> 30) & 0x1FF;
  const uint64_t pd_index = (ivirt >> 21) & 0x1FF;
  const uint64_t pt_index = (ivirt >> 12) & 0x1FF;

  if (!(pml4[pml4_index] & PAGE_PRESENT)) {
    return;
  }

  uint64_t *pdpt = PHYS_TO_VIRT(pml4[pml4_index] & ~(uint64_t)0xFFF);
  if (!(pdpt[pdpt_index] & PAGE_PRESENT)) {
    return;
  }

  uint64_t *pd = PHYS_TO_VIRT(pdpt[pdpt_index] & ~(uint64_t)0xFFF);
  if (!(pd[pd_index] & PAGE_PRESENT)) {
    return;
  }

  uint64_t *pt = PHYS_TO_VIRT(pd[pd_index] & ~(uint64_t)0xFFF);
  memory_page_free((void *)(pt[pt_index] & ~(uint64_t)0xFFF));
  pt[pt_index] = 0;

  __asm__ volatile("invlpg (%0)" : : "r"(virt) : "memory");
}