/* Memory manager functions,
 *
 * Uses a bitmap to keep track of 4k pages that are in use.
 *
 * */

#include "kernel.h"

// Bitmap that keeps track of all the 4 kB memory pages in the system:
static uint64_t *frames = 0;
static uint64_t nframes = 0;

memmap_t* kernel_map = 0; // kernel memory mapping
memmap_t* current_map = 0; // The currently active memory mapping

static void set_frame(uint64_t frame)
{
   uint64_t idx = frame / 64;
   uint64_t off = frame % 64;
   frames[idx] |= (0x1 << off);
}

static void clear_frame(uint64_t frame)
{
   uint64_t idx = frame / 64;
   uint64_t off = frame % 64;
   frames[idx] &= ~(0x1 << off);
}

uint64_t test_frame(uint64_t frame)
{
   uint64_t idx = frame / 64;
   uint64_t off = frame % 64;
   return (frames[idx] & (0x1 << off));
}

static uint64_t first_frame()
{
   uint64_t i, j;
   for (i = 0; i < nframes / 64; i++)
   {
      if (frames[i] != 0xFFFFFFFFFFFFFFFF)
      {
         for (j = 0; j < 64; j++)
         {
            uint64_t dut = 0x1 << j;
            if ((frames[i] & dut) != dut)
            {
               return i*64+j;
            }
         }
      }
   }

   // No frame found:
   return (uint64_t) -1;
}

/*
 * Initializes the memory manager, allocating a bitmap once.
 */
void init_memory(uint64_t total_mem_size)
{
   printf("Running with %d MB ram\n", total_mem_size / 1000000);

   // Only here placement malloc is used!
   //
   // Allocate and clear bits to remember which 4KiB-frames are in use:
   nframes = (total_mem_size / 0x1000); // Calculate number of frames
   frames = (uint64_t*)kmalloc_int( (nframes / 64) * sizeof(uint64_t) );
   memset(frames, 0, (nframes / 64) * sizeof(uint64_t));

   // Create kernel map:
   kernel_map = (memmap_t*)kmalloc_aligned_int(sizeof(memmap_t));
   memset(kernel_map, 0, sizeof(memmap_t));

   // Identity map lower memory and mark as used by the kernel.
   // Mark as used by the kernel:
   uint64_t i;
   i = 0;
   while ( i <= placement_address)
   {
      page_t *page;

      page = get_page(i, kernel_map); // Get page for address
      page->address = i >> 12; // Set the address of this page.
      page->present = 1;
      page->rw = 1;
      page->us = 1; // Make all things accessable for users for now.
      //page->pwt = 1; // Is this bit required?

      set_frame(i / 0x1000);

      i += 0x1000; // Increase a 4 k frame
   }

   // Set the created mapping as active:
   switch_mapping(kernel_map);

   // TODO: set the use of placement malloc to invalid after here.
}

void alloc_frame(page_t *page)
{
   if (page->address != 0)
   {
      return;
   }
   uint64_t idx = first_frame();
   if (idx == (uint64_t) -1)
   {
      panic("No more memory!");
   }
   set_frame(idx);

   page->present = 1;
   page->rw = 1;
   page->us = 0;
   page->address = idx; // set address in frame number, the byte address is 0x1000 times this value.
}

void free_frame(page_t *page)
{
   clear_frame(page->address / 0x1000);
   page->address = 0;
   page->present = 0;
}

void switch_mapping(memmap_t* mapping)
{
   current_map = mapping;

   asm volatile("mov %0, %%cr3" :: "r"(&mapping->table));

   // Enable paging (and flush cache):
   uint64_t cr0;
   asm volatile("mov %%cr0, %0": "=r"(cr0));
   cr0 |= 0x80000000;
   asm volatile("mov %0, %%cr0":: "r"(cr0));
}

/* Get a page for a virtual address, and create other tables if required */
page_t* get_page(uint64_t address, memmap_t *mapping)
{
   // Retrieve the correct PDP (page dir pointer table):
   uint64_t pml4index = (address >> 39) & 0x1FF;
   PDPT_t *pdpt = 0;
   if (mapping->table[pml4index].present == 1)
   {
      pdpt = (PDPT_t*)((uint64_t)mapping->table[pml4index].address << 12);
   }
   else
   {
      // Create new table:
      pdpt = (PDPT_t*)kmalloc_aligned_int(sizeof(PDPT_t));
      memset(pdpt, 0, sizeof(PDPT_t));

      // TODO: get function like virt2phys or something here
      uint64_t address = (uint64_t)pdpt; // get the physical address
      mapping->table[pml4index].address = address >> 12; // shift right
      mapping->table[pml4index].rw = 1;
      mapping->table[pml4index].us = 1;
      mapping->table[pml4index].present = 1;
   }

   uint64_t pdptindex = (address >> 30) & 0x1FF;
   // Retrieve the correct page directory:
   PD_t *pd = 0;
   if (pdpt->table[pdptindex].present == 1)
   {
      pd = (PD_t*)((uint64_t)pdpt->table[pdptindex].address << 12);
   }
   else
   {
      // Create a new table:
      pd = (PD_t*)kmalloc_aligned_int(sizeof(PD_t));
      memset(pd, 0, sizeof(PD_t));

      // Enter table into PDPT:
      // TODO: make virt2phys function:
      uint64_t address = (uint64_t)pd;
      pdpt->table[pdptindex].address = address >> 12;
      pdpt->table[pdptindex].rw = 1;
      pdpt->table[pdptindex].us = 1;
      pdpt->table[pdptindex].present = 1;
   }

   // Retrieve the correct page table:
   uint64_t pdindex = (address >> 21) & 0x1FF;
   PT_t *pt = 0;
   if (pd->table[pdindex].present == 1)
   {
      pt = (PT_t*)((uint64_t)pd->table[pdindex].address << 12);
   }
   else
   {
      // Create table:
      pt = (PT_t*)kmalloc_aligned_int(sizeof(PD_t));
      memset(pt, 0, sizeof(PT_t));

      // Enter PT into PD:
      uint64_t address = (uint64_t)pt;
      pd->table[pdindex].address = address >> 12;
      pd->table[pdindex].rw = 1;
      pd->table[pdindex].us = 1;
      pd->table[pdindex].present = 1;
   }

   // Finally get the page from the directory:
   // TODO: convert from physical address to virtual address:
   uint64_t ptindex = (address >> 12) & 0x1FF;
   return &pt->table[ptindex];
}