#include <mm/vmm.h>
#include <sys/io.h>
#include <ubixos/kpanic.h>
#include <lib/kprintf.h>
#include <lib/kmalloc.h>
#include <ubixos/vitals.h>
#include <ubixos/spinlock.h>
#include <assert.h>

int		numPages = 0;
mMap		*mmUsedPages;

uInt32 countMemory();

void
mmFreeTaskPages(kTask_t *task)
{
	/* just move the list to tail of the free page list */
	if(task->FirstPage != NULL)
	{
		mmFreePages->Last->Next = task->FirstPage;
		task->FirstPage->Previous = mmFreePages->Last;
		mmFreePages->Last = task->LastPage;
		kprintf("mmFreeTaskPages: Memory has been wiped\n");
	}
	else
		kprintf("mmFreeTaskPages: Nothing to free!\n");	
	return;
}


/* TODO: free contiguous blocks of memory */
void 
mmFreeVirtualPage(kTask_t *task, uInt32 memAddr)
{
	mMap *tmp;

	for(tmp = task->FirstPage ; tmp != NULL ; tmp = tmp->Next)
	{
		if(tmp->pageAddr == memAddr)
		{
			/* remove the page from the task's used list */
			if(tmp == task->FirstPage)
			{
				task->FirstPage = task->FirstPage->Next;
			}
			else
			{
				tmp->Previous->Next = tmp->Next;
				tmp->Next->Previous = tmp->Previous;
				tmp->Next = NULL;
				tmp->Previous = NULL;
			}
			
			/* add the page to the end free pages list */
			mmFreePages->Last->Next = tmp;
			tmp->Previous = mmFreePages->Last;
			mmFreePages->Last = tmp;
			kprintf("mmFreeVirtualPage: %d has been freed\n", memAddr);
			return;
		}
	}
	kpanic("mmFreeVirtualPage: attempted to free non-existant page\n");
	return;
}

void *
mmGetFreeVirtualPage(kTask_t *task)
{
	mMap *tmp;

	if(mmFreePages == NULL)
	{
		kprintf("Out of memory\n");
		return NULL;
	}

	/* remove the first page from the list and return it */
	tmp = mmFreePages->First;
	mmFreePages->First = mmFreePages->First->Next;
	mmFreePages->First->Previous = NULL;
	return tmp;
}

int
mmMemMapInit()
{
	uInt32 memStart, z;
	mMap *tmpMap;

	/* Count System Memory */
	numPages = countMemory();

	/* calculate the start of memory */
	memStart  = 0x101;
	memStart += (((sizeof(mMap) * numPages) + (sizeof(mMap) - 1)) / 0x1000);

	/* initialize free pages */
	for(z = memStart ; z < numPages; z++)
	{
		if(mmFreePages == NULL)
		{
			//UBU: replace this with static location
			mmFreePages = kmalloc(sizeof(mMap));
			mmFreePages->First = mmFreePages;
			mmFreePages->Last = mmFreePages;
			mmFreePages->Next = NULL;
			mmFreePages->Previous = NULL;
			mmFreePages->pageAddr = z * 4096;
		}
		else
		{
			//UBU: replace this with static location
			tmpMap = kmalloc(sizeof(mMap));
			mmFreePages->Last->Next = tmpMap;
			tmpMap->Previous = mmFreePages->Last;
			mmFreePages->Last = tmpMap;
			tmpMap->pageAddr = z * 4096;
		}
	}

	/* initialize used pages (kernel space) */
	for(z = 0 ; z < memStart; z++)
	{
		if(mmUsedPages == NULL)
                {
			mmUsedPages = kmalloc(sizeof(mMap));
                        mmUsedPages->First = mmUsedPages;
                        mmUsedPages->Last = mmUsedPages;
                        mmUsedPages->Next = NULL;
                        mmUsedPages->Previous = NULL;
			mmUsedPages->pageAddr = z * 4096;
		}
		else
		{
                        tmpMap = kmalloc(sizeof(mMap));
                        mmUsedPages->Last->Next = tmpMap;
                        tmpMap->Previous = mmUsedPages->Last;
                        mmUsedPages->Last = tmpMap;
			tmpMap->pageAddr = z * 4096;
		}
	}

	/* Print Out Memory Information */
	kprintf("Real Memory:      %iMB\n", ((numPages  * 4096) / 1024) / 1024 );
	kprintf("Available Memory: %iMB\n", ((freePages * 4096) / 1024) / 1024 );
	kprintf("Used Memory:      %iMB\n", ((usedPages * 4096) / 1024) / 1024 );

	/* Return */
	return (0);
}

uInt32 countMemory() {
  register uInt32 *mem = 0x0;
  unsigned long    memCount = -1, tempMemory = 0x0;
  unsigned short   memKb = 0;
  unsigned char    irq1State, irq2State;
  unsigned long    cr0 = 0x0;

  /*
   * Save The States Of Both IRQ 1 And 2 So We Can Turn Them Off And Restore
   * Them Later
   */
  irq1State = inportByte(0x21);
  irq2State = inportByte(0xA1);

  /* Turn Off IRQ 1 And 2 To Prevent Chances Of Faults While Examining Memory */
  outportByte(0x21, 0xFF);
  outportByte(0xA1, 0xFF);

  /* Save The State Of Register CR0 */
  asm volatile (
    "movl %%cr0, %%ebx\n"
    :               "=a" (cr0)
    :
    :               "ebx"
    );

  asm volatile ("wbinvd");
  asm volatile (
    "movl %%ebx, %%cr0\n"
    :
    :               "a" (cr0 | 0x00000001 | 0x40000000 | 0x20000000)
    :               "ebx"
    );

  while (memKb < 4096 && memCount != 0) {
    memKb++;
    if (memCount == -1)
      memCount = 0;
    memCount += 1024 * 1024;
    mem = (uInt32 *)memCount;
    tempMemory = *mem;
    *mem = 0x55AA55AA;
    asm("": : :"memory");
    if (*mem != 0x55AA55AA) {
      memCount = 0;
      }
    else {
      *mem = 0xAA55AA55;
      asm("": : :"memory");
      if (*mem != 0xAA55AA55) {
	memCount = 0;
        }
      }
    asm("": : :"memory");
    *mem = tempMemory;
    }

  asm volatile (
    "movl %%ebx, %%cr0\n"
    :
    :               "a" (cr0)
    :               "ebx"
    );

  /* Restore States For Both IRQ 1 And 2 */
  outportByte(0x21, irq1State);
  outportByte(0xA1, irq2State);

  /* Return Amount Of Memory In Pages */
  return ((memKb * 1024 * 1024) / 4096);
  }