#include <string.h>
#include <stdio.h>

#include <mm/paging.h>
#include <mm/pmm.h>

#define PAGE_DIRECTORY_ENTRIES 1024
#define PAGE_TABLE_ENTRIES     1024
#define PAGE_SIZE              4096

typedef uint32_t page_directory_entry_t;
typedef uint32_t page_table_entry_t;

static page_directory_entry_t* page_directory = (page_directory_entry_t*)0x9C000; /* Must be page-aligned */
static page_table_entry_t* page_tables[PAGE_DIRECTORY_ENTRIES];

extern void _enable_paging_asm(void);

void paging_init(void)
{
  /* Allocate and clear the page directory */
  page_directory = (uint32_t*)alloc_page();
  memset(page_directory, 0, 4096);

  /* Allocate and set up the first identity-mapped page table (0–4MB) */
  uint32_t* first_table = (uint32_t*)alloc_page();
  memset(first_table, 0, 4096);
  for (uint32_t i = 0; i < 1024; i++)
  {
    first_table[i] = (i * 0x1000) | 3; /* Present | RW */
  }
  page_directory[0] = ((uint32_t)first_table) | 3;

  /* --- Map 4MB heap at 0xC0000000 --- */
  uint32_t* heap_table = (uint32_t*)alloc_page();
  memset(heap_table, 0, 4096);

  for (uint32_t i = 0; i < 1024; i++) /* 256 pages = 1MB */
  {
    uint32_t phys = (uint32_t)alloc_page();
    if (phys == 0)
    {
      printf("Out of physical memory during heap mapping!\n");
      while (1);
    }
    heap_table[i] = phys | 3; /* Present | RW */
  }

  /* 0xC0000000 >> 22 == 768 */
  page_directory[768] = ((uint32_t)heap_table) | 3;

  // Load page directory
  asm volatile ("mov %0, %%cr3" : : "r"(page_directory));

  /* Enable paging */
  _enable_paging_asm();
}

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

  /* Allocate page table if necessary */
  if (!(page_directory[pd_index] & 1))
  {
    void* pt_phys = alloc_page();
    page_tables[pd_index] = (page_table_entry_t*)((uint32_t)pt_phys + 0xC0000000); /* Map it higher */
    
    for (int i = 0; i < PAGE_TABLE_ENTRIES; i++)
    {
      page_tables[pd_index][i] = 0x00000002; /* Not present */
    }
    
    page_directory[pd_index] = ((uint32_t)pt_phys) | 0x3; /* Present, R/W */
  }

  page_table_entry_t* page_table = (page_table_entry_t*)((page_directory[pd_index] & 0xFFFFF000) | 0xC0000000);
  page_table[pt_index] = ((uint32_t)phys_addr & 0xFFFFF000) | 0x3; /* Present, R/W */

  asm volatile ("invlpg (%0)" :: "r" (virt_addr) : "memory");
}