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


void 
mmFreeVirtualPage(kTask_t *task, uInt32 memAddr)
{
	mMap *tmp, *tmp2;

	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;

			/* free contiguous blocks */
			if(tmp->Link != NULL)
			{
				kprintf("Freeing contiguous memory\n");
				for(tmp2 = tmp->Link ; tmp2 != NULL ; tmp2 = tmp2->Link)
				{
					mmFreePages->Last->Next = tmp2;
					tmp2->Previous = mmFreePages->Last;
					mmFreePages->Last = tmp2;
				}
			}

			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 + offset
			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 + offset
			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;
	uInt8		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);
}