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

#include "kernel.h"

#define HEAP_START 0xC000000000000000
#define HEAP_INITIAL_SIZE  0x100000

// 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
// TODO: determine if this is required.
memmap_t* current_map = 0; // The currently active memory mapping

heap_t* kernel_heap = 0;

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;
}

void memory_status(void)
{
   // Determine amount of used memory.
   uint64_t i, j;
   uint64_t inuse = 0;
   for (i = 0; i < nframes / 64; i++)
   {
      for (j = 0; j < 64; j++)
      {
         if ((frames[i] & (1UL<<j)) == (1UL<<j))
         {
            inuse++;
         }
      }
   }
   printf("Memory: %d, in use: %d\n", nframes/0x100, inuse/0x100);
}

/*
 * 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);

   // Memory size is from address 0x100000
   total_mem_size += 0x100000;

   // Only here placement malloc is used!
   //
   // Allocate and clear bits to remember which 4KiB-frames are in use:
   // Calculate number of frames. memory size is indicated from address 0x100000
   nframes = (total_mem_size) / 0x1000; // Calculate number of frames
   int numbytes = ((nframes / 64) + 1) * sizeof(uint64_t);
   frames = (uint64_t*)kmalloc(numbytes);
   memset(frames, 0, numbytes);

   // Create kernel map:
   kernel_map = (memmap_t*)kmalloc_a(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 + 0x100000)
   {
      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); // Loads cr3

   // TODO: set the use of placement malloc to invalid after here.
   // kernel_heap = create_heap(HEAP_START, HEAP_INITIAL_SIZE);
}

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;

   // Load table address:
   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:
   printf("Get page for address %x\n", address);
      // TODO: make sure that this function calls kmalloc instead of placement malloc.
      pdpt = (PDPT_t*)kmalloc_a(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_a(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_a(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];
}