#include <mm/pmm.h>
#include <string.h>
#include <mm_macros.h>
#include <stdio.h>

#include <mm/paging.h>

#define PAGE_PRESENT 0x1
#define PAGE_WRITE   0x2
#define PAGE_SIZE    4096

static uint32_t* page_directory;

void map_page(void* phys_addr, void* virt_addr)
{
  uint32_t pd_idx = ((uint32_t)virt_addr >> 22) & 0x3FF;
  uint32_t pt_idx = ((uint32_t)virt_addr >> 12) & 0x3FF;

  uint32_t* page_table;
  
  if (!(page_directory[pd_idx] & PAGE_PRESENT))
  {
    page_table = (uint32_t*) pmm_alloc_page();
    MEMSET(page_table, 0, PAGE_SIZE);
    page_directory[pd_idx] = ((uint32_t)page_table) | PAGE_PRESENT | PAGE_WRITE;
  }
  else
  {
    page_table = (uint32_t*)(page_directory[pd_idx] & ~0xFFF);
  }

  page_table[pt_idx] = ((uint32_t)phys_addr) | PAGE_PRESENT | PAGE_WRITE;
}

void paging_init(void)
{
#ifdef _DEBUG
  printf("[ PAGING ] Initializing paging...\n");
#endif

  page_directory = (uint32_t*)pmm_alloc_page();
  
  MEMSET(page_directory, 0, PAGE_SIZE);

  for (uint32_t addr = 0; addr < 0x800000; addr += PAGE_SIZE)
  {
    map_page((void*) addr, (void*) addr); /* identity map first 8MB */
  }

  asm volatile("mov %0, %%cr3" :: "r"(page_directory));
  uint32_t cr0;
  asm volatile("mov %%cr0, %0" : "=r"(cr0));
  cr0 |= 0x80000000;
  asm volatile("mov %0, %%cr0" :: "r"(cr0));
  
#ifdef _DEBUG
  printf("[ PAGING ] Paging initialized\n");
#endif
}