#include <stdlib.h>

#include "kernel.h"
#include "vmem.h"

#include "asm_x86.h"
#include "mem.h"
#include "kmalloc.h"
#include "lib/string/string.h"

//! i86 architecture defines 1024 entries per table--do not change
#define PAGES_PER_TABLE 1024
#define PAGES_PER_DIR	1024

#define PAGE_DIRECTORY_INDEX(x) (((x) >> 22) & 0x3ff)
#define PAGE_TABLE_INDEX(x) (((x) >> 12) & 0x3ff)
#define PAGE_GET_PHYSICAL_ADDRESS(x) (*x & ~0xfff)

//! page table represents 4mb address space
#define PTABLE_ADDR_SPACE_SIZE 0x400000

//! directory table represents 4gb address space
#define DTABLE_ADDR_SPACE_SIZE 0x100000000

//! page sizes are 4kb
#define PAGE_SIZE 4096

//! page table entry
typedef uint32_t pt_entry;

//! a page directery entry
typedef uint32_t pd_entry;

//! page table
typedef struct ptable_struct {
    pt_entry m_entries[PAGES_PER_TABLE];
}ptable ;
 
//! page directory
typedef struct pdirectory_struct {
    pd_entry m_entries[PAGES_PER_DIR];
}pdirectory;
static uint32_t kernel_pages;
uint32_t fb_addr; // TODO!??!!? how can we share so ugly with vesa
static uint32_t apic_addr;
static uint32_t io_apic_addr;

// TODO : why is the frame not 0xfffff??
enum PAGE_PTE_FLAGS 
{
    I86_PTE_PRESENT		=	1,		//0000000000000000000000000000001
    I86_PTE_WRITABLE		=	2,		//0000000000000000000000000000010
    I86_PTE_USER		=	4,		//0000000000000000000000000000100
    I86_PTE_WRITETHOUGH		=	8,		//0000000000000000000000000001000
    I86_PTE_NOT_CACHEABLE	=	0x10,		//0000000000000000000000000010000
    I86_PTE_ACCESSED		=	0x20,		//0000000000000000000000000100000
    I86_PTE_DIRTY		=	0x40,		//0000000000000000000000001000000
    I86_PTE_PAT			=	0x80,		//0000000000000000000000010000000
    I86_PTE_CPU_GLOBAL		=	0x100,		//0000000000000000000000100000000
    I86_PTE_LV4_GLOBAL		=	0x200,		//0000000000000000000001000000000
    I86_PTE_FRAME		=	0xFFFFF000 	//1111111111111111111000000000000
};

enum PAGE_PDE_FLAGS 
{
    I86_PDE_PRESENT		=	1,		//0000000000000000000000000000001
    I86_PDE_WRITABLE		=	2,		//0000000000000000000000000000010
    I86_PDE_USER		=	4,		//0000000000000000000000000000100
    I86_PDE_PWT			=	8,		//0000000000000000000000000001000
    I86_PDE_PCD			=	0x10,		//0000000000000000000000000010000
    I86_PDE_ACCESSED		=	0x20,		//0000000000000000000000000100000
    I86_PDE_DIRTY		=	0x40,		//0000000000000000000000001000000
    I86_PDE_4MB			=	0x80,		//0000000000000000000000010000000
    I86_PDE_CPU_GLOBAL		=	0x100,		//0000000000000000000000100000000
    I86_PDE_LV4_GLOBAL		=	0x200,		//0000000000000000000001000000000
    I86_PDE_FRAME		=	0xFFFFF000 	//1111111111111111111000000000000
};

////

//! virtual address
typedef uint32_t virtual_addr;
typedef uint32_t physical_addr;
////

static void pt_entry_add_attrib (pt_entry* e, uint32_t attrib)
{
    *e|=attrib;
}
static void pt_entry_del_attrib (pt_entry* e, uint32_t attrib)
{
    *e&=~attrib;
}
static void pt_entry_set_frame (pt_entry* e , physical_addr addr)
{
    *e|=I86_PTE_FRAME&addr;
}

static physical_addr pt_entry_get_frame (pt_entry* e)
{
    return *e&I86_PTE_FRAME;
}
static bool pt_entry_is_present (pt_entry e)
{
    return e&I86_PTE_PRESENT;
}

static bool pt_entry_is_writable (pt_entry e)
{
    return e&I86_PTE_WRITABLE;
}

static physical_addr pt_entry_pfn (pt_entry e)
{
    return e&I86_PTE_FRAME;
}

static void pd_entry_add_attrib (pd_entry* e, uint32_t attrib)
{
    *e|=attrib;
}

static void pd_entry_del_attrib (pd_entry* e, uint32_t attrib)
{
    *e&=~attrib;
}

static void pd_entry_set_frame (pd_entry* e, physical_addr add)
{
    *e|=I86_PDE_FRAME&add;
}

static physical_addr pd_entry_get_frame (pd_entry* e)
{
    return *e&I86_PDE_FRAME;
}

static bool pd_entry_is_present (pd_entry e)
{
    return e&I86_PDE_PRESENT;
}

/*
bool pd_entry_is_user (pd_entry e)
{
    return 1;
}

bool pd_entry_is_4mb (pd_entry e)
{
    return 1;
}
*/

static bool pd_entry_is_writable (pd_entry e)
{
    return e&I86_PDE_WRITABLE;
}

static physical_addr pd_entry_pfn (pd_entry e)
{
    return e&I86_PDE_FRAME;
}

/*
void pd_entry_enable_global (pd_entry e)
{
}
*/

static uint8_t vmmngr_alloc_page (pt_entry* e) 
{
    // allocate a free physical frame
    void* p = mem_alloc_block ();
    if (!p) return 0;

    // map it to the page
    pt_entry_set_frame (e, (physical_addr)p);
    pt_entry_add_attrib (e, I86_PTE_PRESENT);

    return 1;
}

static void vmmngr_free_page (pt_entry* e) 
{
    void* p = (void*)pt_entry_pfn (*e);
    if (p) mem_free_block (p);
    pt_entry_del_attrib (e, I86_PTE_PRESENT);
}

static pt_entry* vmmngr_ptable_lookup_entry (ptable* p, virtual_addr addr) 
{
    if (p) return &p->m_entries[ PAGE_TABLE_INDEX (addr) ];
    return 0;
}

/// SHIT BELOW THIs LINE
void vmem_free_dir(pdirectory *dir)
{
    x86_paging_disable();

    uint32_t virt_addr=0;

    // first pages are identity mapped (kernel space) just remove links
    for(int j=0;j<3;j++)
    {
	    pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	    virt_addr+=1024*4096;
    }

    // programm space 
    virt_addr=0x8000000;
    for(int j=0;j<2;j++)
    {

	pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	ptable *table=pd_entry_get_frame(entry);
    
	for (int i=0, virt=virt_addr; i<1024; i++, virt+=4096) 
	{
	    pt_entry pte=&table->m_entries [PAGE_TABLE_INDEX (virt) ];
	    mem_free_block(pt_entry_get_frame(pte));   // get free space from the memory manager
	}

	mem_free_block(pd_entry_get_frame (entry));
	virt_addr+=1024*4096;

    }

    // stack
    virt_addr=0x8c00000;
    for(int j=0;j<1;j++)
    {

	pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	ptable *table=pd_entry_get_frame(entry);
    
	for (int i=0, virt=virt_addr; i<1024; i++, virt+=4096) 
	{
	    pt_entry pte=&table->m_entries [PAGE_TABLE_INDEX (virt) ];
	    mem_free_block(pt_entry_get_frame(pte));   // get free space from the memory manager
	}

	mem_free_block(pd_entry_get_frame (entry));
	virt_addr+=1024*4096;

    }
    mem_free_block(dir);

    x86_paging_enable();
}



//
//  vmem init /  also copies all the shit over for FORK 
//  if copy_dir==NULL creates brandnew dir 
//  otherwise copies kernel pages and copies 
//  programm pages procreates new programmspace
//
//    TODO: FIX
//   KERNEL SPACE `kernel_pages` first 8 PAGES = 32megs
//   FRAMEBUFER WE GET ON INIT
//   PROGRAMM SPACE HARDCODED TO 0x8000000+2 pages and 0x8c00000+1 pages
//   

pdirectory* vmem_new_space_dir(pdirectory *copy_dir,bool stack_only)
{
    x86_paging_disable();

    //
    pdirectory* dir = (pdirectory*) kballoc(1);

    klog("new pdirectory: 0x%X",dir);
    if (!dir)kpanic("unable to alloc pdirectory");

    // first of all let's zero all the entries
    for(int i=0;i<1024;i++)dir->m_entries [i]=0;

    // identity mapping for first blocks
    uint32_t phys_addr=0;
    uint32_t virt_addr=0;

    kernel_pages=8;

    // first pages are identity mapped
    for(int j=0;j<kernel_pages;j++)
    {

	// this is the table for our page directory. It maps 1024*4096 bytes
	ptable* table;

	// create new tables on init
	if(copy_dir==NULL)
	{
	    // alloc space for our new table
	    table = (ptable*) kballoc(1);
	    if (!table)kpanic("unable to alloc table");

	    //! idenitity mapping
	    for (int i=0, frame=phys_addr, virt=virt_addr; i<1024; i++, frame+=4096, virt+=4096) 
	    {
		//! create a new page
		pt_entry page=0;
		pt_entry_add_attrib (&page, I86_PTE_PRESENT);
		pt_entry_add_attrib (&page, I86_PTE_WRITABLE);
		pt_entry_add_attrib (&page, I86_PTE_USER);
		pt_entry_set_frame (&page, frame);

		//! ...and add it to the page table
		table->m_entries [PAGE_TABLE_INDEX (virt) ] = page;
	    }

	    pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	    *entry=0;
	    pd_entry_add_attrib (entry, I86_PDE_PRESENT);
	    pd_entry_add_attrib (entry, I86_PDE_WRITABLE);
	    pt_entry_add_attrib (entry, I86_PTE_USER);
	    pd_entry_set_frame (entry, (physical_addr)table);

	}

	// otherwise simply take existing stuff from pdir 0
	else
	{
	    dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)]=
		    copy_dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)];

	}

	phys_addr+=1024*4096;
	virt_addr+=1024*4096;
    }

    // FRAMEBUFFER IDENTITY MAP (4pages enought? )
    phys_addr=fb_addr;
    virt_addr=fb_addr;

    if(fb_addr>0x100000)
    for(int j=0;j<4;j++)
    {

	// this is the table for our page directory. It maps 1024*4096 bytes
	ptable* table;

	// create new tables on init
	if(copy_dir==NULL)
	{
	    // alloc space for our new table
	    table = (ptable*) kballoc(1);
	    if (!table)kpanic("unable to alloc table");

	    //! idenitity mapping
	    for (int i=0, frame=phys_addr, virt=virt_addr; i<1024; i++, frame+=4096, virt+=4096) 
	    {
		//! create a new page
		pt_entry page=0;
		pt_entry_add_attrib (&page, I86_PTE_PRESENT);
		pt_entry_add_attrib (&page, I86_PTE_WRITABLE);
		pt_entry_add_attrib (&page, I86_PTE_USER);
		pt_entry_set_frame (&page, frame);

		//! ...and add it to the page table
		table->m_entries [PAGE_TABLE_INDEX (virt) ] = page;
	    }

	    pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	    *entry=0;
	    pd_entry_add_attrib (entry, I86_PDE_PRESENT);
	    pd_entry_add_attrib (entry, I86_PDE_WRITABLE);
	    pt_entry_add_attrib (entry, I86_PTE_USER);
	    pd_entry_set_frame (entry, (physical_addr)table);

	}

	// otherwise simply take existing stuff from pdir 0
	else
	{
	    dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)]=
		    copy_dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)];

	}

	phys_addr+=1024*4096;
	virt_addr+=1024*4096;
    }

    // 2 pages for apic ( in case on boundary??)
    phys_addr=apic_addr;
    virt_addr=apic_addr;

    for(int j=0;j<2;j++)
    {

	// this is the table for our page directory. It maps 1024*4096 bytes
	ptable* table;

	// create new tables on init
	if(copy_dir==NULL)
	{
	    // alloc space for our new table
	    table = (ptable*) kballoc(1);
	    if (!table)kpanic("unable to alloc table");

	    //! idenitity mapping
	    for (int i=0, frame=phys_addr, virt=virt_addr; i<1024; i++, frame+=4096, virt+=4096) 
	    {
		//! create a new page
		pt_entry page=0;
		pt_entry_add_attrib (&page, I86_PTE_PRESENT);
		pt_entry_add_attrib (&page, I86_PTE_WRITABLE);
		pt_entry_add_attrib (&page, I86_PTE_USER);
		pt_entry_set_frame (&page, frame);

		//! ...and add it to the page table
		table->m_entries [PAGE_TABLE_INDEX (virt) ] = page;
	    }

	    pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	    *entry=0;
	    pd_entry_add_attrib (entry, I86_PDE_PRESENT);
	    pd_entry_add_attrib (entry, I86_PDE_WRITABLE);
	    pt_entry_add_attrib (entry, I86_PTE_USER);
	    pd_entry_set_frame (entry, (physical_addr)table);

	}

	// otherwise simply take existing stuff from pdir 0
	else
	{
	    dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)]=
		    copy_dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)];

	}

	phys_addr+=1024*4096;
	virt_addr+=1024*4096;
    }
    // 2 pages for ioapic ( in case on boundary??)
    phys_addr=io_apic_addr;
    virt_addr=io_apic_addr;

    for(int j=0;j<2;j++)
    {

	// this is the table for our page directory. It maps 1024*4096 bytes
	ptable* table;

	// create new tables on init
	if(copy_dir==NULL)
	{
	    // alloc space for our new table
	    table = (ptable*) kballoc(1);
	    if (!table)kpanic("unable to alloc table");

	    //! idenitity mapping
	    for (int i=0, frame=phys_addr, virt=virt_addr; i<1024; i++, frame+=4096, virt+=4096) 
	    {
		//! create a new page
		pt_entry page=0;
		pt_entry_add_attrib (&page, I86_PTE_PRESENT);
		pt_entry_add_attrib (&page, I86_PTE_WRITABLE);
		pt_entry_add_attrib (&page, I86_PTE_USER);
		pt_entry_set_frame (&page, frame);

		//! ...and add it to the page table
		table->m_entries [PAGE_TABLE_INDEX (virt) ] = page;
	    }

	    pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	    *entry=0;
	    pd_entry_add_attrib (entry, I86_PDE_PRESENT);
	    pd_entry_add_attrib (entry, I86_PDE_WRITABLE);
	    pt_entry_add_attrib (entry, I86_PTE_USER);
	    pd_entry_set_frame (entry, (physical_addr)table);

	}

	// otherwise simply take existing stuff from pdir 0
	else
	{
	    dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)]=
		    copy_dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)];

	}

	phys_addr+=1024*4096;
	virt_addr+=1024*4096;
    }
    // programm space 
    virt_addr=0x8000000;
    for(int j=0;j<2;j++)
    {
	if(stack_only)
	{
	    dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)]=
	    copy_dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)];
	    virt_addr+=1024*4096;
	    continue;
	}

	ptable* table = (ptable*) kballoc (1);
	pd_entry *oldentry=NULL;
	ptable* oldtable=NULL;

	if(copy_dir!=NULL)
	{
	    oldentry=&(copy_dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)]);
	    oldtable=pd_entry_get_frame(oldentry);
	}

	klog("oldtable at: 0x%08X",oldtable);

	if (!table)kpanic("unable to alloc table");
	klog("alloc table: %08X",table);
    
	for (int i=0, virt=virt_addr; i<1024; i++, virt+=4096) 
	{
	    //klog("i = %d",i);
	    phys_addr=mem_alloc_block();   // get free space from the memory manager
	    if (!phys_addr)kpanic("unable to alloc spce for frame");

	    // if this is not init , copy contents from current space!
	    if(copy_dir!=NULL)
	    {
		uint32_t addr_old=pt_entry_get_frame(&oldtable->m_entries[PAGE_TABLE_INDEX(virt)]);
		//klog("physical memcpy(0x%08X,0x%08X,4096)",phys_addr, addr_old);
		memcpy(phys_addr,addr_old,4096);

	    }

	    uint32_t frame=phys_addr;

	    //! create a new page
	    pt_entry page=0;
	    pt_entry_add_attrib (&page, I86_PTE_PRESENT);
	    pt_entry_add_attrib (&page, I86_PTE_WRITABLE);
	    pt_entry_add_attrib (&page, I86_PTE_USER);
	    pt_entry_set_frame (&page, frame);

	    //! ...and add it to the page table
	    table->m_entries [PAGE_TABLE_INDEX (virt) ] = page;
	}
	
	pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	*entry=0;
	pd_entry_add_attrib (entry, I86_PDE_PRESENT);
	pd_entry_add_attrib (entry, I86_PDE_WRITABLE);
	    pt_entry_add_attrib (entry, I86_PTE_USER);
	pd_entry_set_frame (entry, (physical_addr)table);

	virt_addr+=1024*4096;
    }

    // programm space  : stack?
    virt_addr=0x8c00000;
    for(int j=0;j<1;j++)
    {

	ptable* table = (ptable*) kballoc (1);
	pd_entry *oldentry=NULL;
	ptable* oldtable=NULL;

	if(copy_dir!=NULL)
	{
	    oldentry=&(copy_dir->m_entries[PAGE_DIRECTORY_INDEX(virt_addr)]);
	    oldtable=pd_entry_get_frame(oldentry);
	}
	klog("oldtable at: 0x%08X",oldtable);

	if (!table)kpanic("unable to alloc table");
	klog("alloc table: %08X",table);
    
	for (int i=0, virt=virt_addr; i<1024; i++, virt+=4096) 
	{
	    phys_addr=mem_alloc_block();   // get free space from the memory manager
	    if (!phys_addr)kpanic("unable to alloc spce for frame");

	    // if this is not init , copy contents from current space!
	    if(copy_dir!=NULL)
	    {
		uint32_t addr_old=pt_entry_get_frame(&oldtable->m_entries[PAGE_TABLE_INDEX(virt)]);
		//klog("physical memcpy(0x%08X,0x%08X,4096)",phys_addr, addr_old);
		memcpy(phys_addr,addr_old,4096);

	    }

	    uint32_t frame=phys_addr;

	    //! create a new page
	    pt_entry page=0;
	    pt_entry_add_attrib (&page, I86_PTE_PRESENT);
	    pt_entry_add_attrib (&page, I86_PTE_WRITABLE);
	    pt_entry_add_attrib (&page, I86_PTE_USER);
	    pt_entry_set_frame (&page, frame);

	    //! ...and add it to the page table
	    table->m_entries [PAGE_TABLE_INDEX (virt) ] = page;
	}
	
	pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt_addr) ];
	*entry=0;
	pd_entry_add_attrib (entry, I86_PDE_PRESENT);
	pd_entry_add_attrib (entry, I86_PDE_WRITABLE);
        pt_entry_add_attrib (entry, I86_PTE_USER);
	pd_entry_set_frame (entry, (physical_addr)table);

	virt_addr+=1024*4096;

    }

    if(copy_dir==NULL) // this happens only on init
    {
	klog("initializing virtual memory (paging)");
	x86_set_page_directory(dir);
    }

    x86_paging_enable();
    return dir;
}

pdirectory* vmem_init(uint32_t kernel_blocks, uint32_t frameb_addr, uint32_t apic_ad, uint32_t ioapic_ad)
{
    fixme("do not share fb_addr with syscalls like that!");
    fb_addr=frameb_addr;
    apic_addr=apic_ad;
    io_apic_addr=ioapic_ad;
    kernel_pages=kernel_blocks/1024+1;
    return vmem_new_space_dir(NULL,false);
}