/*
 * This file contains all functions used to manipulate the virtual address
 * spaces.  It has a static Kernel page directory and table, which it uses to
 * initialsed an identity-paged environment for the Kernel to work in.  This is
 * enough for the heap to function in.  The file also exposes several functions
 * that allow a page directory to be manipulated and have pages added and moved.
 * These functions are used by other components of the Kernel - mostly the heap
 * and IPC.  There are also functions to create new and destroy existing page
 * directories.  The paging system also implements features like copy-on-write.
 */

#include <stdint.h>
#include <sys/mman.h>
#include "frame.h"
#include "heap.h"
#include "mem.h"
#include "paging.h"
#include "../vfs/vfs.h"
#include "../proc/proc.h"
#include "../io.h"
#include "../screen.h"

Spinlock quickPageLock;

page_table_t kernelDir;
page_t zeroFrame;

void enable_paging(void);
void disable_paging(void);
void copy_page_frame(void *src, void *dest);

/* Switch page directory */
static void
switch_dir(page_dir_t dir)
{
	asm volatile("mov %0, %%cr3" :: "r" (dir));
}

/* Allocate a page a frame */
void
alloc_page(page_t *page, uint16_t flags, page_t frame)
{
	page_t *mappings = (void *) 0xFFC00000;
	page_table_t *tables = (void *) 0xFFFFF000;
	if ((tables[(page - mappings) / 1024] & PDE_PRESENT) == 0)
		return;
	if (*page & 0xFFFFF000)
		return;

	if (frame == (page_t) -1)
		frame = alloc_frames(1);
	if (frame == (page_t) -1)
		return;

	*page = frame | flags;
	flush_tlb((page - mappings) << 12);
}

/* Release a page's frame */
void
free_page(page_t *page)
{
	page_t *mappings = (void *) 0xFFC00000;
	page_table_t *tables = (void *) 0xFFFFF000;
	if ((tables[(page - mappings) / 1024] & PDE_PRESENT) == 0)
		return;
	if ((*page & 0xFFFFF000) == 0)
		return;

	free_frame(*page & 0xFFFFF000);
	*page = 0x00000000;
	flush_tlb((page - mappings) << 12);
}

/* Get Page Table Entry from virtual address */
page_t *
get_page(void *addr)
{
	page_t *mappings = (void *) 0xFFC00000;
	page_table_t *tables = (void *) 0xFFFFF000;
	uint32_t address = (uint32_t) addr >> 12;
	uint32_t tbl = address / 1024;
	/* Create table not present */
	if ((tables[tbl] & PDE_PRESENT) == 0) {
		tables[tbl] = alloc_frames(1)
		            | PDE_PRESENT | PDE_WRITE | PDE_USER;
		memset((void *) mappings + (tbl * 0x1000), 0, 0x1000);
	}
	return &mappings[address];
}

/* Clone a page directory */
page_dir_t
clone_dir(void)
{
	page_table_t *oldTables = (void *) 0xFFFFF000;
	page_table_t *newTables = (void *) 0xFFFFE000;
	page_t *oldTable, *newTable;
	page_dir_t dir = alloc_frames(1);
	uint16_t i, tbl, pg;

	/* Temporarily link new paging structures into current directory */
	page_table_t restore = oldTables[1022];
	oldTables[1022] = dir | PDE_PRESENT | PDE_WRITE;
	for (i = 0; i < 1024; i++)
		flush_tlb((uintptr_t) newTables + (0x1000 * i));

	/* Iterate tables */
	for (tbl = 0; tbl < 1022; tbl++) {
		if ((oldTables[tbl] & PDE_PRESENT) == 0)
			continue;

		/* Link Kernel tables */
		if (tbl < 2 || tbl >= 1008) { /* TODO: define kernel mem */
			newTables[tbl] = oldTables[tbl];
			continue;
		}

		/* Copy everything else */
		newTables[tbl] = alloc_frames(1) | PG_ATTR(oldTables[tbl]);
		oldTable = (page_t *) 0xFFC00000 + (tbl * 1024);
		newTable = (page_t *) 0xFF800000 + (tbl * 1024);
		for (pg = 0; pg < 1024; pg++) {
			if ((oldTable[pg] & PTE_PRESENT) == 0) {
				newTable[pg] = 0;
				continue;
			}

			/* Copy-On-Write behaviour */
			if (tbl < 960) {
				oldTable[pg] &= ~PTE_WRITE;
				flush_tlb((uintptr_t) (((tbl * 1024) + pg) << 12));
				newTable[pg] = oldTable[pg];
			} else {
				newTable[pg] = alloc_frames(1) | PG_ATTR(oldTable[pg]);
				copy_page_frame((void *) PG_ADDR(oldTable[pg]),
				                (void *) PG_ADDR(newTable[pg]));
			}
			/* FIXME */
		}
	}
	newTables[1023] = oldTables[1022];

	/* Unlink paging structures */
	oldTables[1022] = restore;
	for (i = 0; i < 1024; i++)
		flush_tlb((uintptr_t) newTables + (0x1000 * i));

	return dir;
}

/* Free all (copied) pages in the current directory */
void
clean_dir(void)
{
	page_t *mappings = (void *) 0xFFC00000;
	page_table_t *tables = (void *) 0xFFFFF000;
	page_t *pages;
	uint16_t tbl, pg;
	for (tbl = 2; tbl < 1008; tbl++) {
		if ((tables[tbl] & PDE_PRESENT) == 0)
			continue;
		pages = mappings + (tbl * 1024);
		for (pg = 0; pg < 1024; pg++) {
			if ((pages[pg] & PDE_PRESENT) == 0)
				continue;
			free_page(pages + pg);
		}
	}
}

/* Quickly map a page frame into view for temporary use */
page_t
quick_page(uintptr_t frame)
{
	page_t *mappings = (void *) 0xFFC00000;
	page_t old;
	old = mappings[2047];
	mappings[2047] = PG_ADDR(frame) | PG_ATTR(old);
	flush_tlb(0x7FF000);
	return PG_ADDR(old);
}

/* Initialise paging */
void
init_paging(void)
{
	zeroFrame = alloc_frames(1);
	memset((void *) zeroFrame, 0, 0x1000);

	uint16_t tbl, pg;
	page_t *table;
	page_table_t *kernelTables;
	kernelDir = alloc_frames(1);
	kernelTables = (page_table_t *) kernelDir;
	for (tbl = 0; tbl < 1024; tbl++)
		kernelTables[tbl] = 0x00000000 | PDE_WRITE;
	for (tbl = 0; tbl < 2; tbl++) {
		table = (void *) alloc_frames(1);
		kernelTables[tbl] = ((page_table_t) table)
		                  | PDE_WRITE | PDE_PRESENT;
		for (pg = 0; pg < 1024; pg++) {
			if (!tbl && !pg)
				continue;
			table[pg] = (((tbl * 1024) + pg) << 12)
			          | PTE_WRITE | PTE_PRESENT | PTE_GLOBAL;
		}
	}
	/* Map the directory into itself */
	kernelTables[1023] = kernelDir | PDE_WRITE | PDE_PRESENT;

	/* Use Kernel directory */
	switch_dir(kernelDir);
	register_exception(14, early_page_fault_handler);
	enable_paging();

	/* Identity page the APIC registers */
	*get_page((void *)  lapicPtr) =  lapicPtr
	                              | PTE_PRESENT | PTE_WRITE | PTE_GLOBAL;
	*get_page((void *) ioapicPtr) = ioapicPtr
	                              | PTE_PRESENT | PTE_WRITE | PTE_GLOBAL;

	/* Allocate Kernel stack */
	uintptr_t stk;
	for (stk = 0xF0400000; stk < 0xF0800000; stk += 0x1000)
		alloc_page(get_page((void *) stk),
		           PTE_PRESENT | PTE_WRITE | PTE_USER, -1);
}