#include <asm/system.h>
#include <kernel/memory.h>
#include <stdint.h>
#include <sys/types.h>

static char __attribute__((aligned (4096))) mapped_page[4096];

static pdir_t kpdir;

ptab_t ktab;

static ptab_t flow;
static ptab_t fupp;

void *alloc_physical_page(void) {
  int i;
  void *ret;

  for(i = 0; i < PGENT; i++) {
    if(flow[i] & 1) {
      ret = (void*) (flow[i] & ~1);
      flow[i] = 0;
      return ret;
    }

    if(fupp[i] & 1) {
      ret = (void*) (fupp[i] & ~1);
      fupp[i] = 0;
      return ret;
    }
  }

  return NULL;
}

int free_physical_page(void *page) {
  int i;
  page = (void*) ((uint32_t) page & ~0xFFF);

  for(int i = 0; i < PGENT; i++) {
    if(!(flow[i] & 1)) {
      flow[i] = (uint32_t) page | 1;
      return 0;
    }

    if(!(fupp[i] & 1)) {
      fupp[i] = (uint32_t) page | 1;
      return 0;
    }
  }

  return -1;
}

void *alloc_physical_pages(unsigned int n) {
  void *p = NULL;
  void *last = NULL;
  void **page;

  while(n--) {
    p = alloc_physical_page();
    if(p == NULL) {
      page = &last;
      while(*page) {
	free_physical_page(*page);
        page = *page;
      }
      return NULL;
    }

    map_page(p);
    *((void**) mapped_page) = last;
    last = p;
  }

  return p;
}

void *map_page(void *page) {
  uint32_t p = (uint32_t) mapped_page;

  /* enforce 4KB page alignment */
  page = (void*) ((uint32_t) page & ~0xFFF);

  ktab[(p >> 12) & 0x3FF] = (uint32_t) page | 0x003;

  invlpg(mapped_page);

  return mapped_page;
}

void paging_init(void) {
  int i;

  /* empty the page directory */
  for(i = 0; i < PGENT; i++)
    kpdir[i] = 0;

  /* populate the free page tables */
  for(i = 0; i < PGENT; i++) {
    /* populate the lower table */
    flow[i] = (i << 12);
    /* only pages >1MB can be allocated */
    if(i >= 0x100)
      flow[i] |= 0x001;
    /* populate the upper table */
    fupp[i] = ((i << 12) + 0x400000) | 0x001;
  }

  /* populate the kernel's page table */
  for(i = 0; i < PGENT; i++) {
    ktab[i] = 0;
    if(i < 256)
      ktab[i] = (i << 12) | 0x003;
  }

  /* populate the kernel's page directory */
  kpdir[0] = ((pde_t) ktab) | 0x007;

  __asm__ (
      "mov %%eax, %%cr3\n"
      "mov %%cr0, %%eax\n"
      "or $0x80000000, %%eax\n"
      "mov %%eax, %%cr0\n"
      :: "a" (&kpdir)
      );
}