#include <stdlib.h>

#include "kernel.h"
#include "log.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

static uint32_t e1000_addr;

//! 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 fb_addr;
static uint32_t local_apic_addr;
static uint32_t io_apic_addr;

static uint32_t mod_start;
static uint32_t mod_end;

// 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 pt_entry_is_user (pt_entry e)
{
    return e&I86_PTE_USER;
}

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


static pdirectory*  vmem_clean_dir()
{
    fixme("replace ALLL kballoc with getting the pages from mem/vmem");
    pdirectory* dir = (pdirectory*) kballoc(1);
    for(int i=0;i<1024;i++)dir->m_entries [i]=0; // zero all
    return dir;
}

static void vmem_clear_one(pdirectory* dir,uint32_t virt)
{
	pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt) ];
	ptable* table;
	if(*entry!=0){
	    table= pd_entry_get_frame (entry);
	    table->m_entries [PAGE_TABLE_INDEX (virt) ] = 0;
	}
}


// addresses need to be page aligned. (or will be forced down)
static void vmem_del_generic(pdirectory* dir, uint32_t virt,uint32_t pages, bool dealloc, bool user)
{
    //fixme("make sure the pages are marked as used in the physical mem manager, really?");
    
    //force align
    virt/=4096;
    virt*=4096;

    while(pages>0)
    {
	pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt) ];
	ptable* table;

	if(*entry==0){
	    // no pagetable here yet.
            kpanic("no pagetable found for vmem entry to delete!");
	}
	else{
	    // we use the existing one.
	    table= pd_entry_get_frame (entry);
	}

	if(dealloc)
	{
	    uint32_t page=table->m_entries [PAGE_TABLE_INDEX (virt) ];
	    uint32_t phys=pt_entry_get_frame(page);
            mem_free_block(phys); // free the space
	}

	//zero  page
	pt_entry page=0;

	//set the page table
	uint32_t old=table->m_entries [PAGE_TABLE_INDEX (virt) ];

	if(old==0)kpanic("nothing to delete in vmem at 0x%08x",virt);

	table->m_entries [PAGE_TABLE_INDEX (virt) ] = page;

	virt+=4096;

	pages--;
    }
}

// addresses need to be page aligned. (or will be forced down)
static void vmem_add_generic(pdirectory* dir,uint32_t phys,uint32_t virt,uint32_t pages, bool alloc, bool user)
{
    //fixme("make sure the pages are marked as used in the physical mem manager, really?");

    //force align
    phys/=4096;
    phys*=4096;

    //force align
    virt/=4096;
    virt*=4096;

    while(pages>0)
    {
	pd_entry* entry = &dir->m_entries [PAGE_DIRECTORY_INDEX (virt) ];
	ptable* table;

	if(*entry==0){
	    // no pagetable here yet.
	    table = (ptable*) kballoc(1);
	    for(int i=0;i<1024;i++)table->m_entries [i]=0; // zero all

	    pd_entry_add_attrib (entry, I86_PDE_PRESENT);
	    pd_entry_add_attrib (entry, I86_PDE_WRITABLE);
	    if(user)pd_entry_add_attrib (entry, I86_PDE_USER);
	    pd_entry_set_frame (entry, table);
	}
	else{
	    // add to existing one.
	    table= pd_entry_get_frame (entry);
	}

	if(alloc)
	{
	    phys=mem_alloc_block(); // get free space from the memory manager
//	    klog("allocated physical at 0x%08X",phys);
	}

	//create a new page
	pt_entry page=0;
	pt_entry_add_attrib (&page, I86_PTE_PRESENT);
	pt_entry_add_attrib (&page, I86_PTE_WRITABLE);
	if(user)pt_entry_add_attrib (&page, I86_PTE_USER);
	pt_entry_set_frame  (&page, phys);

	//and add it to the page table
	uint32_t old=table->m_entries [PAGE_TABLE_INDEX (virt) ];
	if(old!=0&&old!=page)kpanic("overwriting existing page data! virt=0x%08X phys=0x%08X",virt,phys);
	table->m_entries [PAGE_TABLE_INDEX (virt) ] = page;

	virt+=4096;
	phys+=4096;

	pages--;
    }
}

static void vmem_add_alloc(pdirectory* dir,uint32_t addr,uint32_t pages,bool user)
{
    vmem_add_generic(dir,addr,addr,pages, true,user);
}

static void vmem_add_remap(pdirectory* dir,uint32_t phys,uint32_t virt,uint32_t pages,bool user)
{
    vmem_add_generic(dir,phys,virt,pages, false,user);
}

static void vmem_add_identity(pdirectory* dir,uint32_t addr,uint32_t pages, bool user)
{
    vmem_add_generic(dir,addr,addr,pages, false,user);
}

pdirectory* vmem_kernel_dir() // we create thie once for each cpu
{
    fixme("remove user flags where appropriate");
    fixme("do not waste soo many pages/page tables!");
    fixme("align properly!! / merge page tables if required!");
    fixme("Is ioapic/lapic really on one page?");
    fixme("kernelspace: kernel/lapic/ioapic/framebuffer/ext2-ramimage/stack/cpu private");
    fixme("check ext2 and kernel and framebufffer do not exceen the 32megs)");

    pdirectory* dir = vmem_clean_dir();

    vmem_add_identity(dir,VMEM_KERNEL,VMEM_KERNEL_PAGES,false);//identity map first 32 megs... 
    vmem_add_identity(dir,e1000_addr,32,false);//identity map 32 pages for e1000

    vmem_add_remap(dir,fb_addr,VMEM_FRAMEBUFFER,VMEM_FRAMEBUFFER_PAGES,false);//32megs should be enough for 4k (think about pitch)

    vmem_add_remap(dir,local_apic_addr,VMEM_LAPIC,1,false); //apic addr should be at pagestart, right? TODO: check.
    vmem_add_remap(dir,io_apic_addr,VMEM_IOAPIC,1,false); 
    vmem_add_remap(dir,mod_start,VMEM_EXT2_RAMIMAGE,VMEM_EXT2_PAGES,false);//32megs for ramimage: TODO: check if enough? 

    vmem_add_alloc(dir,VMEM_CPU_PRIVATE,4,false);
    vmem_add_alloc(dir,VMEM_CPU_STACK_TOP-4096*VMEM_CPU_STACK_PAGES,VMEM_CPU_STACK_PAGES,false);

    vmem_add_alloc(dir,VMEM_USER_PROG,VMEM_USER_PROG_PAGES,true); 
    vmem_add_alloc(dir,VMEM_USER_ENV,1,true); 
    vmem_add_alloc(dir,VMEM_USER_NEWLIB,1,true); 
    vmem_add_alloc(dir,VMEM_USER_STACK_TOP-4096*VMEM_USER_STACK_PAGES,VMEM_USER_STACK_PAGES,true); 
    //vmem_add_alloc(dir,VMEM_USER_FRAMEBUFFER,VMEM_USER_FRAMEBUFFER_PAGES,true); /// each new process gets a 640x480x32 area

    return dir;
}

void vmem_add_framebuffer(pdirectory *dir)
{
//    vmem_del_generic(dir,VMEM_USER_FRAMEBUFFER,VMEM_USER_FRAMEBUFFER_PAGES,false, true);
    vmem_add_alloc(dir,  VMEM_USER_FRAMEBUFFER,VMEM_USER_FRAMEBUFFER_PAGES,true); /// each new process gets a 640x480x32 area
}

void vmem_add_total_framebuffer(pdirectory *dir)
{
    vmem_add_remap(dir,fb_addr,VMEM_FRAMEBUFFER,VMEM_FRAMEBUFFER_PAGES,true);//32megs should be enough for 4k (think about pitch)
}

void vmem_free_space_dir(pdirectory *dir,bool stack_only)
{

    fixme("free kmalloc too!");
    fixme("stack_only version too!");
    klog("free page directory at: 0x%X",dir);

    if(stack_only)kpanic("not implemented yet!");

    //free user pages
    uint32_t virt=0;
    for(int i=0;i<1024;i++)
    {
	uint32_t src_pt=dir->m_entries [i];

	if(pt_entry_is_user(src_pt))
	{
	    ptable *src_table=pt_entry_get_frame(&src_pt);

	    for(int j=0;j<1024;j++)
	    {
		uint32_t src_pd=src_table->m_entries[j];
	        uint32_t src_phys=pd_entry_get_frame(&src_pd);

		if(src_pd)
		{
                    if(virt>=VMEM_USER_FRAMEBUFFER&&virt<VMEM_USER_FRAMEBUFFER+VMEM_USER_FRAMEBUFFER_PAGES*4096)
                    {
                    }
                    else
                    {
		    mem_free_block(src_phys);
                    }
		}

		virt+=4096;
	    }
	}
	else virt+=4096*1024;
    }

}

pdirectory* vmem_new_space_dir(pdirectory *copy_dir,bool stack_only)
{
    pdirectory* dir = vmem_clean_dir(); //let's start with a fresh page directory
    klog("new page directory at: 0x%X",dir);

    //link non-user pages.
    for(int i=0;i<1024;i++)
    {
	uint32_t src_pt=copy_dir->m_entries [i];
	if(!pt_entry_is_user(src_pt))dir->m_entries [i]=src_pt;
    }

    if(!stack_only)
    {
     // temp disable //   vmem_add_alloc(dir,VMEM_USER_FRAMEBUFFER,VMEM_USER_FRAMEBUFFER_PAGES,true); /// each new process gets a 640x480x32 area
	vmem_add_alloc(dir,VMEM_USER_PROG,VMEM_USER_PROG_PAGES,true); 
	vmem_add_alloc(dir,VMEM_USER_ENV,1,true); 
    }

    // threads need own //
    vmem_add_alloc(dir,VMEM_USER_NEWLIB,1,true); 
    vmem_add_alloc(dir,VMEM_USER_STACK_TOP-4096*VMEM_USER_STACK_PAGES,VMEM_USER_STACK_PAGES,true); 

    pdirectory* mydir=x86_get_page_directory();

    //copy user pages
    uint32_t virt=0;
    for(int i=0;i<1024;i++)
    {
	uint32_t src_pt=copy_dir->m_entries [i];
	uint32_t dst_pt=dir->m_entries [i];

	if(pt_entry_is_user(src_pt))
	{
	    ptable *src_table=pt_entry_get_frame(&src_pt);
	    ptable *dst_table=pt_entry_get_frame(&dst_pt);

	    for(int j=0;j<1024;j++)
	    {
		uint32_t src_pd=src_table->m_entries[j];
		uint32_t dst_pd=dst_table->m_entries[j];

	        uint32_t src_phys=pd_entry_get_frame(&src_pd);
	        uint32_t dst_phys=pd_entry_get_frame(&dst_pd);

		if(src_pd)
		{
                    if(stack_only&&
                            (
                             virt==VMEM_USER_ENV
                             ||(virt>=VMEM_USER_PROG&&virt<VMEM_USER_PROG+VMEM_USER_PROG_PAGES*4096)
//                             ||(virt>=VMEM_USER_FRAMEBUFFER&&virt<VMEM_USER_FRAMEBUFFER+VMEM_USER_FRAMEBUFFER_PAGES*4096)
                            )
                      )
                    {
                        // threads share this //
                        vmem_add_remap(dir,src_phys,virt,1,true);
                    }
                    else
                    {
                             if(virt>=VMEM_USER_FRAMEBUFFER&&virt<VMEM_USER_FRAMEBUFFER+VMEM_USER_FRAMEBUFFER_PAGES*4096)
                            {
                                // and for now we all share this! //
                                // or not!
                                //vmem_add_remap(dir,src_phys,virt,1,true);
                            }
                             else{
                    
                    
                        //klog("copy virt: %x / phys: %x -> %x",virt,src_phys,dst_phys);
                    
                        vmem_clear_one(mydir,VMEM_COPY_PAGE);
                        vmem_clear_one(mydir,VMEM_COPY_PAGE+4096);

                        vmem_add_remap(mydir,src_phys,VMEM_COPY_PAGE,1,false);
                        vmem_add_remap(mydir,dst_phys,VMEM_COPY_PAGE+4096,1,false);

                        x86_invlpg(VMEM_COPY_PAGE); // refresh TLB 
                        x86_invlpg(VMEM_COPY_PAGE+4096); // refresh TLB 

                        memcpy(VMEM_COPY_PAGE+4096,VMEM_COPY_PAGE,4096);
                             }
                    }
                }
		virt+=4096;
	    }
	}
	else virt+=4096*1024;
    }

    return dir;
}

void vmem_init(multiboot_information *cfg_multiboot, acpi_information *cfg_acpi,uint32_t e1000addr)
{
    e1000_addr=e1000addr;
    fixme("do not share fb_addr with syscalls like that!");
    fb_addr=cfg_multiboot->framebuffer_addr;
    local_apic_addr=cfg_acpi->local_apic_address;
    io_apic_addr=cfg_acpi->io_apic_address;
    multiboot_mod *mod=(multiboot_mod *)cfg_multiboot->mods_addr;
    mod_start=mod->mod_start;
    mod_end=mod->mod_end;
}