#include "mem.h"
#include "i386.h"
#include "multiboot.h"
#include "lib.h"

struct GDT {
  u32_t base;
  u32_t limit;
  u8_t type;
};

struct segdesc_s gdt[GDT_SIZE];

struct tss_s tss;

struct gdt_ptr gdtr;

void encodeGdtEntry(unsigned char target[8], struct GDT source)
{
  if (source.limit > 65536) {
    // Adjust granularity if required
    source.limit = source.limit >> 12;
    target[6] = 0xC0;
  } else {
    target[6] = 0x40;
  }
    
  target[0] = source.limit & 0xFF;
  target[1] = (source.limit >> 8) & 0xFF;
  target[6] |= (source.limit >> 16) & 0xF;
    
  target[2] = source.base & 0xFF;
  target[3] = (source.base >> 8) & 0xFF;
  target[4] = (source.base >> 16) & 0xFF;
  target[7] = (source.base >> 24) & 0xFF;
    
  target[5] = source.type;
}

/* Set up GDT */
void seg_init() {
  struct GDT tmp = {0, 0, 0};
  
  // null
  encodeGdtEntry((u8_t *)gdt, tmp);

  // code segment
  tmp.limit = 0xffffffff;
  tmp.type = 0x9a;
  encodeGdtEntry((u8_t *)(gdt+CS_INDEX), tmp);

  // data segment
  tmp.type = 0x92;
  encodeGdtEntry((u8_t *)(gdt+DS_INDEX), tmp);

  // tss
  tss.ss0 = DS_SELECTOR;
  tmp.base = (u32_t) &tss;
  tmp.limit = sizeof(tss) - 1;
  tmp.type = 0x89;
  encodeGdtEntry((u8_t *)(gdt+TSS_INDEX), tmp);

  // TODO: IDT, LDT

  // lgdt and reload registers
  gdtr.base = gdt;
  gdtr.limit = sizeof(gdt) - 1;
  setGdt();
}

/* Page-Directory-Pointer-Table, which contains 4 Page-Directory-Entries */
static u64_t page_dir_ptr_tab[4] __attribute__((aligned(0x20))); // must be aligned to (at least)0x20

/* Set up page table, with PAE enabled */
void page_init() {
  u32_t cr0;

  unsigned int i;

  /* We will start the paging data structures  after the kernel stack */

  /* Page-Directory-Table, which contains 512 entries */
  u64_t *page_dir = (u64_t *) STACK_SIZE;  // Takes one page: 512 entries * 8 (bytes/entry)
  // zero page directory table
  for(i = 0; i < VM_DIR_ENTRIES; i++)
    {
      //attribute: supervisor level, read/write, not present.
      page_dir[i] = 0 | 2; 
    }  
  page_dir_ptr_tab[0] = (u32_t) page_dir | VM_PRESENT;

  u64_t *page_table = page_dir + VM_DIR_ENTRIES;

  /* Fill in some entries */
  // TODO: create a mmap function to manage the mappings
  // TODO: read memory info provided by grub

  // map the first 512 pages, i.e. 2MB
  page_dir[0] = (u32_t) page_table | VM_WRITE | VM_PRESENT ;
  unsigned long address=0; // holds the physical address of where a page is
  for(i = 0; i < 512; i++)
    {
      page_table[i] = address | VM_WRITE | VM_PRESENT;
      address = address + PAGE_SIZE;
    };
  page_table += VM_DIR_ENTRIES;

  // the first page table has been filled up, map the next 512 pages
  page_dir[1] = (u32_t) page_table | VM_WRITE | VM_PRESENT ;
  for(i = 0; i < 512; i++)
    {
      page_table[i] = address | VM_WRITE | VM_PRESENT;
      address = address + PAGE_SIZE;
    };
  page_table += VM_DIR_ENTRIES;

  // For testing purposes, we map 0xc0000000 (the 3rd PDPT entry) to 0
  // Just reuse the other page tables
  page_dir_ptr_tab[3] = page_dir_ptr_tab[0];

  // fast hack, should be moved to i386.S
  //asm volatile("movl %cr4, %eax; bts $5, %eax; movl %eax, %cr4"); // set bit5 in CR4 to enable PAE
  asm volatile("movl %cr4, %eax; orl $0x20, %eax; movl %eax, %cr4"); // set bit5 in CR4 to enable PAE 
  asm volatile ("movl %%eax, %%cr3" :: "a" (&page_dir_ptr_tab)); // load PDPT into CR3
  asm volatile ("movl %cr0, %eax; orl $0x80000000, %eax; movl %eax, %cr0;"); // activate paging

  u32_t BASE = 0xc0000000;
  // NOTE: We don't want the compiler to optimize!!
  volatile u32_t *ptr = (u32_t *) (3*1024*1024);
  *ptr = 0xDEADBEAF;
  ptr = (u32_t *) (3*1024*1024 + BASE);
  printf ("Value at 0x%x: 0x%x\n", ptr, *ptr);
}