// implement fork from user space

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

#undef debug

#define PTE_COW		0x800

Pte *uvpt = (Pte *) UVPT;
Pde *uvpd = (Pde *) UVPT + (UVPT >> PGSHIFT);

extern void _pgfault_entry(void);

//
// Custom page fault handler - if faulting page is copy-on-write,
// map in our own private writable copy.
//

static void
pgfault(u_int va, u_int err)
{
        va = va - (va%BY2PG);
#ifdef debug
        printf("pgfault: va=%08lx, err=%x\n", va, err);
#endif
        
	int r;
        int perm;
	u_char *tmp = NULL;

        if (!(err & FEC_WR))
                panic("pgfault handler called for non-write fault");

        perm = uvpt[va >> PGSHIFT] & (PTE_U | PTE_W | PTE_P | PTE_AVAIL);
        if (!(perm & PTE_COW))
                panic("pgfault handler called for non-CoW page");

        // Find a free page for our temporary page
        int i, j;
        for (i = 0; i < (USTACKTOP >> PGSHIFT); i++) {
                Pde pde = uvpd[i];
                if (pde != 0) {
                        for (j = 0; j < PTE2PT; j++) {
                                if (uvpt[(i << 10) + j] == 0) {
                                        tmp = (u_char *) (i << PDSHIFT)
                                                + (j << PGSHIFT);
                                        goto found;
                                }
                        }
                } else {
                        tmp = (u_char *) (i << PDSHIFT);
                        goto found;
                }
        }
        panic("couldn't find a free page in the CoW pgfault handler");
        
found:
#ifdef debug
        printf("found temporary page, va=%08lx\n", tmp);
#endif
        
        if ((r = sys_mem_alloc(0, (u_long) tmp, PTE_U | PTE_W | PTE_P)) != 0)
                panic("sys_mem_alloc failed: %e", r);

        memcpy(tmp, (char *)va, BY2PG);

        perm &= ~PTE_COW;
        perm |= PTE_W;
        
        if ((r = sys_mem_map(0, (u_long) tmp, 0, va, perm)) != 0)
                panic("sys_mem_map failed: %e", r);

        if ((r = sys_mem_unmap(0, (u_long) tmp)) != 0)
                panic("sys_mem_unmap failed: %e", r);

#ifdef debug
        printf("finished copying page\n");
#endif
  
}

//
// Map our virtual page pn (address pn*BY2PG) 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?)
// 
static void
duppage(u_int envid, u_int pn, u_int sforking)
{
	int r;
	u_int addr;
	Pte pte;

        pte = uvpt[pn];
        int perm = pte & (PTE_U | PTE_W | PTE_P | PTE_AVAIL);

        if (sforking) {
                u_long esp = read_esp;
                if (esp < (pn + 1) * BY2PG) {
                        // Stack page, copy-on-write it
                        if (perm & PTE_W || perm & PTE_COW) {
                                perm &= ~PTE_W;
                                perm |= PTE_COW;
                        }
                } else {
                        // Shared page
                }
        } else {
                if (perm & PTE_W || perm & PTE_COW) {
                        perm &= ~PTE_W;
                        perm |= PTE_COW;
                }
        }
        
        addr = pn * BY2PG;

#ifdef debug
        printf("duppage: %08x from env %x to env %x: pte %08lx\n", addr, env->env_id, envid, pte);
#endif

        if ((r = sys_mem_map(0, addr, envid, addr, perm)) != 0)
                panic("sys_mem_map: %e", r);

        if ((r = sys_mem_map(envid, addr, 0, addr, perm)) != 0)
                panic("sys_mem_map: %e", r);
}

//
// User-level fork.  Create a child and then copy our address space
// and page fault handler setup to the child.
//
// Hint: use vpd, vpt, and duppage.
// Hint: remember to fix "env" in the child process!
// 
int
fork(void)
{
        int childid;
        int r;
        
        set_pgfault_handler(pgfault);
                
        childid = sys_env_alloc();
        if (childid == 0) {
                // we're the child
                env = &envs[ENVX(sys_getenvid())];
                return 0;
        }

        // We're the parent, set stuff up.

        // Copy the page directory mapping
        int i, j;
        for (i = 0; i < (UTOP >> PDSHIFT); i++) {
                Pde pde = uvpd[i];
                if (pde != 0) {
#ifdef debug
                        printf("duplicating page table %lx\n", i);
#endif
                        for (j = 0; j < PTE2PT; j++) {
                                if ((((i << 10) + j) << 12) >= USTACKTOP)
                                        break;
                                
                                if (uvpt[(i << 10) + j] != 0)
                                        duppage(childid, (i << 10) + j, 0);
                        }
                } else {

                }
        }
        
        // Map the exception stack
        r = sys_mem_alloc(childid, UXSTACKTOP-BY2PG, PTE_U | PTE_P | PTE_W );
        if (r != 0)
                panic("sys_mem_alloc: %e\n", r);

        r = sys_set_pgfault_entry(childid, (u_long) _pgfault_entry);
        if (r != 0)
                printf("sys_set_pgfault_entry: %e\n", r);

        r = sys_set_status(childid, ENV_RUNNABLE);
        if (r != 0)
                printf("sys_set_status: %e\n", r);

#ifdef debug
                        printf("fork completed %x -> %x\n", env->env_id, childid);
#endif

        
        return childid;
}

// Challenge!
int
sfork(void)
{
        int childid;
        int r;
        
        set_pgfault_handler(pgfault);
                
        childid = sys_env_alloc();
        if (childid == 0) {
                // we're the child
                env = &envs[ENVX(sys_getenvid())];
                return 0;
        }

        // We're the parent, set stuff up.

        // Copy the page directory mapping
        int i, j;
        for (i = 0; i < (UTOP >> PDSHIFT); i++) {
                Pde pde = uvpd[i];
                if (pde != 0) {
#ifdef debug
                        printf("duplicating page table %lx\n", i);
#endif
                        for (j = 0; j < PTE2PT; j++) {
                                if ((((i << 10) + j) << 12) >= USTACKTOP)
                                        break;
                                
                                if (uvpt[(i << 10) + j] != 0)
                                        duppage(childid, (i << 10) + j, 1);
                        }
                } else {

                }
        }
        
        // Map the exception stack
        r = sys_mem_alloc(childid, UXSTACKTOP-BY2PG, PTE_U | PTE_P | PTE_W );
        if (r != 0)
                panic("sys_mem_alloc: %e\n", r);

        r = sys_set_pgfault_entry(childid, (u_long) _pgfault_entry);
        if (r != 0)
                printf("sys_set_pgfault_entry: %e\n", r);

        r = sys_set_status(childid, ENV_RUNNABLE);
        if (r != 0)
                printf("sys_set_status: %e\n", r);

#ifdef debug
                        printf("fork completed %x -> %x\n", env->env_id, childid);
#endif

        
        return childid;
}