// implement fork from user space

#include <inc/string.h>
#include <inc/lib.h>

// PTE_COW marks copy-on-write page table entries.
// It is one of the bits explicitly allocated to user processes (PTE_AVAIL).
#define PTE_COW		0x800

//
// Custom page fault handler - if faulting page is copy-on-write,
// map in our own private writable copy.
//
static void
pgfault(struct UTrapframe *utf)
{
	//cprintf("trapped page fault at %08x\n",utf->utf_fault_va);
    void *addr = (void *) utf->utf_fault_va;
    uintptr_t addr_page = (uintptr_t) ROUNDDOWN(addr, PGSIZE);
	uint32_t err = utf->utf_err;
	int r;

	// Check that the faulting access was (1) a write, and (2) to a
	// copy-on-write page.  If not, panic.
	// Hint:
	//   Use the read-only page table mappings at vpt
	//   (see <inc/memlayout.h>).

	// LAB 4: Your code here.
    //check err
    //cprintf("user mode page fault at %08x with error %08x\n",utf->utf_fault_va,utf->utf_err);
            
    if ((err&FEC_WR) == 0){
        panic("[%08x] this is not a write, error code is: %08x",env->env_id,err);
    }
    //check permissions
    pte_t pte = vpt[addr_page/PGSIZE];
    if ((pte & PTE_COW) == 0)
        panic("page is not copy-on-write: %08x",addr);

	// Allocate a new page, map it at a temporary location (PFTEMP),
	// copy the data from the old page to the new page, then move the new
	// page to the old page's address.
	// Hint:
	//   You should make three system calls.
	//   No need to explicitly delete the old page's mapping.
	
	// LAB 4: Your code here.
	if (sys_page_alloc(0, (void *)PFTEMP, (PTE_W|PTE_U|PTE_P)) < 0) 
        panic("Could not allocate new page");
    memmove(((void *) PFTEMP), ((void *) addr_page), PGSIZE);
    if (sys_page_map(0, ((void *) PFTEMP), 0, ((void *) addr_page), (PTE_U|PTE_W|PTE_P)) < 0)
        panic("Could not remap page");
}


//
// Map our virtual page pn (address pn*PGSIZE) into the target envid
// at the same virtual address.  If the page is writable or copy-on-write,
// the new mapping must be created copy-on-write, and then our mapping must be
// marked copy-on-write as well.  (Exercise: Why mark ours copy-on-write again
// if it was already copy-on-write?)
//
// Returns: 0 on success, < 0 on error.
// It is also OK to panic on error.
// 
static int
duppage(envid_t envid, unsigned pn)
{
	int r;
	void *addr;
	pte_t pte;

	// LAB 4: Your code here.
	addr = (void *) (pn*PGSIZE);
    pte = vpt[pn];
    
    if ((vpt[pn] & PTE_SHARE) != 0) {
		if ((r = sys_page_map(0, addr, envid, addr, (vpt[pn]&PTE_USER))) < 0) {
			panic("sys_page_map: %e", r);
		}
		return 0;	
	}

    if ((pte&(PTE_COW|PTE_W)) != 0) {
        r = sys_page_map(0, addr, envid, addr, PTE_U|PTE_COW|PTE_P);
        if (r < 0)
            return r;
        //cprintf("1 mapping addr cow = %08x\n",addr);
        r = sys_page_map(0, addr, 0, addr, PTE_U|PTE_COW|PTE_P);
        //cprintf("p4\n");
        if (r < 0)
            return r;
        //cprintf("mapping addr cow = %08x\n",addr);
        
    } else {
        //cprintf("mapping addr read only = %08x\n",addr);
        r = sys_page_map(0, addr, envid, addr, PTE_U|PTE_P);
        if (r < 0)
            return r;
    }
	return 0;
}

//
// User-level fork with copy-on-write.
// Set up our page fault handler appropriately.
// Create a child.
// Copy our address space and page fault handler setup to the child.
// Then mark the child as runnable and return.
//
// Returns: child's envid to the parent, 0 to the child, < 0 on error.
// It is also OK to panic on error.
//
// Hint:
//   Use vpd, vpt, and duppage.
//   Remember to fix "env" and the user exception stack in the child process.
//   Neither user exception stack should ever be marked copy-on-write,
//   so you must allocate a new page for the child's user exception stack.
//
envid_t
fork(void)
{
    int r;
    set_pgfault_handler(&pgfault);
	// LAB 4: Your code here.
	envid_t new_env = sys_exofork();
    
    if (new_env < 0)
        panic("Could not create new environment");

    //child process
    if (new_env == 0){
        env = &envs[ENVX(sys_getenvid())];
        return 0;
}

    //copy mappings
    int i,j,pn;
    pde_t pde;
    pte_t pte;
    for(i=(UTEXT>>PDXSHIFT); i < (UTOP>>PDXSHIFT); i++){
        pde = vpd[i];
        if ((pde&PTE_P)!= 0){
            for (j = 0; j < NPTENTRIES; j++){
                pn = (i<<(PDXSHIFT-PTXSHIFT)) + j;
                pte = vpt[pn];
                if (((pte&PTE_P)!=0) && (pn != ((UXSTACKTOP-PGSIZE) >> PTXSHIFT))) {
                    r = duppage(new_env,pn);
                    if (r < 0)
                        panic("could not copy mapping %d from parent to child: %e",i,r);
                }
            }
        }
    }

    //allocate user exception stack
    r = sys_page_alloc(new_env, ((void *) (UXSTACKTOP-PGSIZE)), PTE_P|PTE_U|PTE_W);
    if (r < 0)
        panic("could not allocate page for user exception stack: %e",r);
    //cprintf("page fault handler: %08x\n",&pgfault);
    sys_env_set_pgfault_upcall(new_env, env->env_pgfault_upcall);
    sys_env_set_status(new_env, ENV_RUNNABLE);
    
	return new_env;
}

envid_t
checkpoint(envid_t envid)
{
    set_pgfault_handler(&pgfault);
	sys_env_set_pgfault_upcall(envid, env->env_pgfault_upcall);
    envid_t new_env = sys_env_copy(envid);
	int r;
    
    if (new_env < 0) 
        panic("Could not copy environment");
    
    r = sys_env_copy_mem(envid, new_env, PTE_COW, PTE_SHARE);
    if (r < 0)
        panic("Could not copy memory");
    
    //allocate user exception stack for both environments
    r = sys_page_alloc(envid, ((void *) (UXSTACKTOP-PGSIZE)), PTE_P|PTE_U|PTE_W);
    if (r < 0)
        panic("could not allocate page for user exception stack: %e",r);
    r = sys_page_alloc(new_env, ((void *) (UXSTACKTOP-PGSIZE)), PTE_P|PTE_U|PTE_W);
    if (r < 0)
        panic("could not allocate page for user exception stack: %e",r);
    
    return new_env;
}
// Challenge!
int
sfork(void)
{
	panic("sfork not implemented");
	return -E_INVAL;
}