// 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)
{
	void *addr = (void *) utf->utf_fault_va;
	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 uvpt
	//   (see <inc/memlayout.h>).
	if(!((err & FEC_WR) && (uvpt[(uintptr_t) addr >> PGSHIFT] & PTE_COW)))
		panic("page fault (addr %p)! %c", addr, (err & FEC_WR) ? 'w' : 'r');

	// 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.
	void* temp_addr = (void*) PFTEMP;
	void* fault_addr = ROUNDDOWN(addr, PGSIZE);
	if(sys_page_alloc(0, temp_addr, PTE_P | PTE_W | PTE_U) < 0)
		panic("failed to allocate new page");

	memcpy(temp_addr, fault_addr, PGSIZE);
	sys_page_map(0, temp_addr, 0, fault_addr, PTE_P | PTE_U | PTE_W);
	sys_page_unmap(0, temp_addr);
}

//
// 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 do we need to mark ours
// copy-on-write again if it was already copy-on-write at the beginning of
// this function?)
//
// 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;
	bool change_own = false;
	pte_t new_pte = uvpt[pn];
	pte_t perms = new_pte & (PTE_P | PTE_U | PTE_W | PTE_AVAIL);
	void* addr = (void*) (pn * PGSIZE);

	// If we're writable, remove write permission
	if((new_pte & PTE_W) || (new_pte & PTE_COW)) {
		perms = (perms & ~PTE_W) | PTE_COW;
		change_own = true;
	}

	// Map either with the same permissions or with COW.
	if((r = sys_page_map(0, addr, envid, addr, perms)) < 0)
		return r;

	// Update our own permissions if necessary
	if(change_own) {
		if((r = sys_page_map(0, addr, 0, addr, perms)) < 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 uvpd, uvpt, and duppage.
//   Remember to fix "thisenv" 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)
{
	set_pgfault_handler(pgfault);

	int return_code;
	envid_t forked;

	forked = sys_exofork();
	if(forked < 0) return forked;
	if(forked == 0) { thisenv = &envs[ENVX(sys_getenvid())]; return 0; }

	// Map all accessible page directory entries
	for(int pde_i = 0; pde_i < PDX(UTOP); pde_i++) {
		pde_t pde = uvpd[pde_i];
		if(!(pde & PTE_P)) continue;

		// For each PDE, map all the underlying PTEs
		for(int pte_i = 0; pte_i < NPTENTRIES; pte_i++) {
			int pn = pde_i * NPTENTRIES + pte_i;
			pte_t pte = uvpt[pn];
			if(!(pte & PTE_P)) continue;

			// Do not map user exception stack, though
			if(pn == ((UXSTACKTOP - PGSIZE) >> PGSHIFT)) continue;

			if((return_code = duppage(forked, pn)) < 0) return return_code;
		}
	}

	// Allocate new page for the exception stack
	return_code = sys_page_alloc(forked, (void*) UXSTACKTOP - PGSIZE,
			PTE_P | PTE_U | PTE_W);
	if(return_code < 0) return return_code;

	// Set the upcall entry point
	sys_env_set_pgfault_upcall(forked, thisenv->env_pgfault_upcall);
	sys_env_set_status(forked, ENV_RUNNABLE);

	return forked;
}

// Challenge!
int
sfork(void)
{
	panic("sfork not implemented");
	return -E_INVAL;
}