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

#define UTEMP2USTACK(addr)	((void*) (addr) + (USTACKTOP - PGSIZE) - UTEMP)
#define UTEMP2			(UTEMP + PGSIZE)
#define UTEMP3			(UTEMP2 + PGSIZE)


static int min(int a, int b) {
	return a > b ? b : a;
}
// Helper functions for spawn.
static int init_stack(envid_t child, const char **argv, uintptr_t *init_esp);
static int copy_shared_pages(envid_t child);

// Spawn a child process from a program image loaded from the file system.
// prog: the pathname of the program to run.
// argv: pointer to null-terminated array of pointers to strings,
// 	 which will be passed to the child as its command-line arguments.
// Returns child envid on success, < 0 on failure.
int
spawn(const char *prog, const char **argv)
{
	unsigned char elf_buf[512];
	struct Trapframe child_tf;
	envid_t child, my_id;
	int r, i;
	int prog_fd;
	uintptr_t init_esp;
	uint32_t addr;

	// Insert your code, following approximately this procedure:
	//
	//   - Open the program file.
	//
	//   - Read the ELF header, as you have before, and sanity check its
	//     magic number.  (Check out your load_icode!)
	//
	//   - Use sys_exofork() to create a new environment.
	//
	//   - Set child_tf to an initial struct Trapframe for the child.
	//     Hint: The sys_exofork() system call has already created
	//     a good basis, in envs[ENVX(child)].env_tf.
	//     Hint: You must do something with the program's entry point.
	//     What?  (See load_icode!)
	//
	//   - Call the init_stack() function above to set up
	//     the initial stack page for the child environment.
	//
	//   - Map all of the program's segments that are of p_type
	//     ELF_PROG_LOAD into the new environment's address space.
	//     Use the p_flags field in the Proghdr for each segment
	//     to determine how to map the segment:
	//
	//	* If the ELF flags do not include ELF_PROG_FLAG_WRITE,
	//	  then the segment contains text and read-only data.
	//	  Use read_map() to read the contents of this segment,
	//	  and map the pages it returns directly into the child
	//        so that multiple instances of the same program
	//	  will share the same copy of the program text.
	//        Be sure to map the program text read-only in the child.
	//        Read_map is like read but returns a pointer to the data in
	//        *blk rather than copying the data into another buffer.
	//
	//	* If the ELF segment flags DO include ELF_PROG_FLAG_WRITE,
	//	  then the segment contains read/write data and bss.
	//	  As with load_icode() in Lab 3, such an ELF segment
	//	  occupies p_memsz bytes in memory, but only the FIRST
	//	  p_filesz bytes of the segment are actually loaded
	//	  from the executable file - you must clear the rest to zero.
	//        For each page to be mapped for a read/write segment,
	//        allocate a page in the parent temporarily at UTEMP,
	//        read() the appropriate portion of the file into that page
	//	  and/or use memset() to zero non-loaded portions.
	//	  (You can avoid calling memset(), if you like, if
	//	  page_alloc() returns zeroed pages already.)
	//        Then insert the page mapping into the child.
	//        Look at init_stack() for inspiration.
	//        Be sure you understand why you can't use read_map() here.
	//
	//     Note: None of the segment addresses or lengths above
	//     are guaranteed to be page-aligned, so you must deal with
	//     these non-page-aligned values appropriately.
	//     The ELF linker does, however, guarantee that no two segments
	//     will overlap on the same page; and it guarantees that
	//     PGOFF(ph->p_offset) == PGOFF(ph->p_va).
	//
	//   - Call sys_env_set_trapframe(child, &child_tf) to set up the
	//     correct initial eip and esp values in the child.
	//
	//   - Start the child process running with sys_env_set_status().

	// LAB 5: Your code here.
	
	if ((my_id = sys_getenvid()) < 0) {
		cprintf("problem with getting its own id\n");
		return -1;
	}

	

	if ((prog_fd = open(prog, O_RDONLY)) < 0) {
		cprintf(" problem with opening the program file\n");
		return -1;
	}

	struct Proghdr *ph, *eph;

	if ((r = read(prog_fd, elf_buf, sizeof elf_buf -1)) < 0) {
		cprintf(" problem with reading the program file\n");
		return -1;
	}

	ph = (struct Proghdr *) ((uint8_t *) elf_buf + ((struct Elf *)elf_buf)->e_phoff);

	if (((struct Elf *)elf_buf)->e_magic != ELF_MAGIC) {
        	cprintf("the elf header does not have the magic number \n");
		return -1;
	}

	if ((child = sys_exofork()) < 0) {
		cprintf("problem with exofork\n");
		return child;
	}

	if  ((r = init_stack(child, argv, &init_esp)) < 0) {
		cprintf("problem with initstack\n");
		return r;
	}	

	child_tf = envs[ENVX(child)].env_tf;
	child_tf.tf_eip = ((struct Elf *)elf_buf)->e_entry & 0xFFFFFF;
	child_tf.tf_esp = init_esp;	


	eph = ph + ((struct Elf *)elf_buf)->e_phnum;

	for (; ph < eph; ph++) {
						
		if (ph->p_type == ELF_PROG_LOAD) {
					
			if ((ph->p_flags & ELF_PROG_FLAG_WRITE) == 0) {		
			
				uint32_t * blk;
				
				
				if ( (r = read_map(prog_fd, ph->p_offset, (void **)&blk)) < 0) {
					cprintf("problem with read map \n");
					return -1;
				}
				
				int pages_to_map = ROUNDUP(ph->p_memsz + PGOFF(ph->p_va) - 1, PGSIZE) / PGSIZE; 
				uint32_t initial_page_blk = ROUNDDOWN((uint32_t)blk, PGSIZE);		
				uint32_t initial_page_va = ROUNDDOWN(ph->p_va, PGSIZE);
				uint32_t start_in_file = ROUNDDOWN(ph->p_offset, PGSIZE);
				
			
				for (i = 0; i < pages_to_map; i++ ) {
					if (lazy_map) {
						if ( (r = read_map(prog_fd, start_in_file + PGSIZE*i, (void **)&blk)) < 0) {
							cprintf("problem with read map \n");
							return -1;
						}
					}
					if (sys_page_map(my_id, (void *) (initial_page_blk+ PGSIZE*i), 
							 child, (void *) (initial_page_va + PGSIZE*i), 
							 PTE_P | PTE_U ) < 0) {
						cprintf("problem with mapping the page %d \n", i);
						return -1;
					}		
				}
		
			} else {
				
				int pages_to_map = ROUNDUP(PGOFF(ph->p_va) + ph->p_memsz -1,PGSIZE)/PGSIZE;
				int offset_in_fd = ph->p_offset;
				int offset_in_page = PGOFF(ph->p_va);
				int size_copied = 0;
				int va_to_map_to = ROUNDDOWN(ph->p_va, PGSIZE);
				
				for (i = 0; i < pages_to_map; i++) {
					
					if ((size_copied < ph->p_filesz) && (seek(prog_fd, offset_in_fd) < 0)) {
						cprintf("problem when seeking in the elf file\n");
						return -1;
					}						 

					if (sys_page_alloc(my_id, (void *)UTEMP, PTE_U|PTE_P|PTE_W) < 0) {
						cprintf("could not allocate a page at UTEMP \n");
						return -1;
					}
		
					memset((void *)UTEMP, 0, BLKSIZE);

					if ((size_copied < ph->p_filesz) && 
					    (read(prog_fd, (char *) (UTEMP + offset_in_page), 
						  min(PGSIZE-offset_in_page, ph->p_filesz-size_copied)) < 0)) {
						cprintf("problem when reading from the offset \n");
						return -1;
					}
				
					if (sys_page_map(my_id, (void *) UTEMP, 
						 child, (void *) va_to_map_to, 
						 PTE_U | PTE_W | PTE_P) < 0) {
					cprintf("problem when inserting the page in the kid\n");
					return -1;
					}
					
					offset_in_fd += PGSIZE;
					size_copied += PGSIZE - offset_in_page;
					offset_in_page = 0;
					va_to_map_to += PGSIZE;				
				}				
			}
		}
	}

	if (sys_env_set_trapframe(child, &child_tf) < 0) {
		cprintf("sys set trapframe is not set up \n");
		return -1;
	}

//+
	//set specific pages PTE_SHARE
	for (addr = UTEXT; addr < USTACKTOP; addr += PGSIZE) {
		//check if page table is present
		if ((vpd[PDX((uint32_t)addr)] & PTE_P) == 0) {
			continue;
		}
		//check if page is mapped
		if ((vpt[(uint32_t)addr/PGSIZE] & PTE_P) == 0) {
			continue;
		}

		if ((vpt[(uint32_t)addr/PGSIZE] & PTE_SHARE) == PTE_SHARE) {
			if ((r = sys_page_map(0, (void *)addr, child, (void *)addr, vpt[(uint32_t)addr/PGSIZE] & (PTE_P|PTE_U|PTE_W|PTE_AVAIL))) < 0) {
				panic("sys_page_map: %e", r);
			}
			continue;	
		}
	}
//-


	if (sys_env_set_status(child, ENV_RUNNABLE) < 0) {
		cprintf("problem with setting the environment runnable\n");
	}

	return child;

}



// Spawn, taking command-line arguments array directly on the stack.
int
spawnl(const char *prog, const char *arg0, ...)
{
	return spawn(prog, &arg0);
}


// Set up the initial stack page for the new child process with envid 'child'
// using the arguments array pointed to by 'argv',
// which is a null-terminated array of pointers to null-terminated strings.
//
// On success, returns 0 and sets *init_esp
// to the initial stack pointer with which the child should start.
// Returns < 0 on failure.
static int
init_stack(envid_t child, const char **argv, uintptr_t *init_esp)
{
	size_t string_size;
	int argc, i, r;
	char *string_store;
	uintptr_t *argv_store;

	// Count the number of arguments (argc)
	// and the total amount of space needed for strings (string_size).
	string_size = 0;
	for (argc = 0; argv[argc] != 0; argc++)
		string_size += strlen(argv[argc]) + 1;

	// Determine where to place the strings and the argv array.
	// Set up pointers into the temporary page 'UTEMP'; we'll map a page
	// there later, then remap that page into the child environment
	// at (USTACKTOP - PGSIZE).
	// strings is the topmost thing on the stack.
	string_store = (char*) UTEMP + PGSIZE - string_size;
	// argv is below that.  There's one argument pointer per argument, plus
	// a null pointer.
	argv_store = (uintptr_t*) (ROUNDDOWN(string_store, 4) - 4 * (argc + 1));
	
	// Make sure that argv, strings, and the 2 words that hold 'argc'
	// and 'argv' themselves will all fit in a single stack page.
	if ((void*) (argv_store - 2) < (void*) UTEMP)
		return -E_NO_MEM;

	// Allocate the single stack page at UTEMP.
	if ((r = sys_page_alloc(0, (void*) UTEMP, PTE_P|PTE_U|PTE_W)) < 0)
		return r;

	// Replace this with your code to:
	//
	//	* Initialize 'argv_store[i]' to point to argument string i,
	//	  for all 0 <= i < argc.
	//	  Also, copy the argument strings from 'argv' into the
	//	  newly-allocated stack page.
	//	  Hint: Copy the argument strings into string_store.
	//	  Hint: Make sure that argv_store uses addresses valid in the
	//	  CHILD'S environment!  The string_store variable itself
	//	  points into page UTEMP, but the child environment will have
	//	  this page mapped at USTACKTOP - PGSIZE.  Check out the
	//	  UTEMP2USTACK macro defined above.
	//
	//	* Set 'argv_store[argc]' to 0 to null-terminate the args array.
	//
	//	* Push two more words onto the child's stack below 'args',
	//	  containing the argc and argv parameters to be passed
	//	  to the child's umain() function.
	//	  argv should be below argc on the stack.
	//	  (Again, argv should use an address valid in the child's
	//	  environment.)
	//
	//	* Set *init_esp to the initial stack pointer for the child,
	//	  (Again, use an address valid in the child's environment.)
	//
	// LAB 5: Your code here.
	char * next_string_ptr = string_store;

	for (i = 0; i < argc; i ++){
		argv_store[i] = UTEMP2USTACK(next_string_ptr);
		next_string_ptr +=  strlen(argv[i]) + 1;
	}
	argv_store[argc] = 0;

	next_string_ptr = string_store;
	for (i = 0; i < argc; i ++){
		memmove(next_string_ptr, argv[i], strlen(argv[i])); 
		next_string_ptr += strlen(argv[i]);
		*next_string_ptr = 0;
		next_string_ptr += 1; 
	}

	*(argv_store-1) = UTEMP2USTACK(argv_store);
	*(argv_store-2) = argc;

	*init_esp = UTEMP2USTACK(argv_store - 2);	

	// After completing the stack, map it into the child's address space
	// and unmap it from ours!
	if ((r = sys_page_map(0, UTEMP, child, (void*) (USTACKTOP - PGSIZE), PTE_P | PTE_U | PTE_W)) < 0)
		goto error;
	if ((r = sys_page_unmap(0, UTEMP)) < 0)
		goto error;

	return 0;

error:
	sys_page_unmap(0, UTEMP);
	return r;
}