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

#include <mm/paging.h>
#include <mm/pmm.h>
#include <mm/heap.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 = NULL; /* Will be allocated */
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 = (page_directory_entry_t*)alloc_page();
  memset(page_directory, 0, PAGE_SIZE);

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

  /* Allocate and clear the heap page table */
  uint32_t heap_pd_index = HEAP_START >> 22;  /* 0xC0000000 >> 22 = 768 */
  page_tables[heap_pd_index] = (page_table_entry_t*)alloc_page();
  memset(page_tables[heap_pd_index], 0, PAGE_SIZE);
  
  /* Map 4MB heap pages */
  for (uint32_t i = 0; i < PAGE_TABLE_ENTRIES; i++) /* 1024 pages = 4MB */
  {
    void* phys = alloc_page();
    if (phys == 0)
    {
      printf("Out of physical memory during heap mapping!\n");
      while (1);
    }
    page_tables[heap_pd_index][i] = ((uint32_t)phys & 0xFFFFF000) | 3; /* Present | RW */
  }
  page_directory[heap_pd_index] = ((uint32_t)page_tables[heap_pd_index]) | 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 */
    
    memset(page_tables[pd_index], 0, PAGE_SIZE);
    
    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");
}