#include <types.h>
#include <proc_syscalls.h>
#include <syscall.h>
#include <lib.h>
#include <limits.h>
#include <types.h>
#include <kern/errno.h>
#include <machine/trapframe.h>
#include <thread.h>
#include <synch.h>
#include <addrspace.h>
#include <proc.h>
#include <current.h>
#include <spl.h>
#include <pid.h>
#include <kern/wait.h>
#include <copyinout.h>
#include <vfs.h>
#include <kern/fcntl.h>

static void thread_fork_start(void *tf, unsigned long vmspace);

/*
 * Implementation of fork system call
 * arguments:
 * tf: trapframe to copy
 * err: pointer to error return value
 */
pid_t
sys_fork(struct trapframe *tf, int* err)
{
    *err = 0;
    struct trapframe *new_tf;
    struct addrspace *addr_space;

    // Clone the current process
    struct proc *new_proc = proc_clone(err);
    if (*err) {
        *err = ENOMEM;
        return -1;
    }

    // Allocate for new trapframe
    new_tf = kmalloc(sizeof(struct trapframe));
    if(new_tf == NULL)
    {
        *err = ENOMEM;
        return -1;
    }

    // Copy original tf to new tf
    memmove(new_tf, tf, sizeof(struct trapframe));

    // Allocate for address space
    addr_space = kmalloc(sizeof(struct addrspace));
    if (addr_space == NULL) {
        kfree(new_tf);
        *err = ENOMEM;
        return -1;
    }

    // Copy over address spaces
    struct addrspace *old_as = proc_getas();
	int return_as = as_copy(old_as, &addr_space);
	if (return_as) {
        kfree(new_tf);
        kfree(addr_space);
        *err = ENOMEM;
		return -1;
	}
    
    // Fork thread into created process
    int return_thread = thread_fork("Child Fork", new_proc,
			        thread_fork_start, (void *) new_tf, (unsigned long) addr_space);
    if (return_thread) {
        kfree(new_tf);
        as_destroy(addr_space);
        *err = ENOMEM;
        return -1;
    }

    return new_proc->pid;
}

/*
 * Initial instructions to newly forked process
 */
static void
thread_fork_start(void *tf, unsigned long vmspace)
{
    int spl = splhigh();
    
    // Unused for now
    (void)vmspace;

    // Create trapframe on local stack, free the malloc'ed one
	struct trapframe child_frame = *(struct trapframe*) tf;
    kfree(tf);

	// change the register values
	child_frame.tf_v0 = 0;
	child_frame.tf_a3 = 0;
	child_frame.tf_epc += 4;

	mips_usermode(&child_frame);

    splx(spl);

	return;
}

/*
 * Implementation of getpid system call
 */
pid_t
sys_getpid(void)
{
    pid_t pid = curproc->pid;
    return pid;
}

/*
 * Kernel-Accessibly helper function for the sys_waitpid system call
 */
int
wait_pid(pid_t pid, int* status, int options)
{
    if(pid == curproc->pid)
    {
        return EINVAL;
    }

    if (options != 0 && options != WNOHANG)
    {
        return EINVAL;
    }
    if (status == NULL)
    {
        return EFAULT;
    }

    if(pid <= INVALID_PID)
    {
        return ENOSYS;
    }

    lock_acquire(pid_list_lock);

    pid_t parent_pid = curproc->pid;

    // Check that process is child before waiting
    if(!pid_is_child(pid, parent_pid))
    {
        lock_release(pid_list_lock);
        return ECHILD;
    }

    struct pid_node * parent = pid_get_node(curproc->pid);

    if(parent == NULL)
    {
        lock_release(pid_list_lock);
        return ESRCH;
    }

    struct pid_node *child = pid_get_node(pid);

    if(child == NULL)
    {
        lock_release(pid_list_lock);
        return EINVAL;
    }

    // wait on condition variable if child is not yet finished
    if(child->finished == false)
    {
        cv_wait(child->pid_cv, pid_list_lock);
        KASSERT(child->finished);
    }

    *status = child->exit_status;

    pid_remove_child(pid, curproc->pid);

    lock_release(pid_list_lock);

    pid_remove(pid);

    return 0;
}

/*
 * Implementation of wait_pid system call
 * arguments:
 * pid: process to wait on to exit
 * status: process exit status
 * options: options to vary waitpid behaviour
 * err: pointer to error return value
 */
pid_t
sys_waitpid(pid_t pid, userptr_t status, int options, int* err)
{
    *err = 0;
    int return_status;

    if (pid > __PID_MAX || pid < 2) {
        *err = ESRCH;
        return -1;
    }

    if (options != 0) {
        *err = EINVAL;
        return -1;
    }

    // bulk of work done in kernel-accessible wait_pid function
    *err = wait_pid(pid, &return_status, options);
    if(*err)
    {
        return -1;
    }

    if(status != NULL)
    {
        *err = copyout(&return_status, status, sizeof(int));
    }

	return pid;
}

/*
 * Implementation of exit system call
 * arguments:
 * status: exit code to report
 */
__DEAD void
sys_exit(int status)
{
    lock_acquire(pid_list_lock);
    
    struct proc *proc_to_exit = curproc;
    KASSERT(proc_to_exit != kproc);

    struct pid_node *cur_pid = pid_get_node(proc_to_exit->pid);
    KASSERT(cur_pid != NULL);
    
    cur_pid->finished = true;
    cur_pid->exit_status = status;

    // Dissociate Children
    while(cur_pid->child_list != NULL)
    {
        struct pid_node *child_pid = pid_get_node(cur_pid->child_list->pid);

        if(child_pid == NULL)
        {
            break;
        }

        if(child_pid->finished)
        {
            pid_remove_child(child_pid->pid, cur_pid->pid);
            pid_remove(child_pid->pid);
        }
        cur_pid->child_list = cur_pid->child_list->next;
    }

    if(cur_pid->parent_pid != INVALID_PID) // Has parent
    {
        cv_broadcast(cur_pid->pid_cv, pid_list_lock);
    }
    else // Has no parent
    {
        lock_release(pid_list_lock);
        pid_remove(proc_to_exit->pid);
    }

    lock_release(pid_list_lock);

    // remove thread from exiting proc
    KASSERT(curthread->t_proc == proc_to_exit);
    proc_remthread(curthread);
    proc_addthread(kproc, curthread);
    KASSERT(threadarray_num(&proc_to_exit->p_threads) == 0);

    proc_destroy(proc_to_exit);

    // exit process thread
    thread_exit();

    // exit kernel thread
    thread_exit();
}

/*
 * Implementation of execv system call
 * arguments:
 * program: program to replace currently executing program
 * args: arguments to pass to program
 * err: pointer to error return value
 */
int
sys_execv(userptr_t program, userptr_t args, int* err) {
    *err = 0;

    if (program == NULL || args == NULL) {
        *err = EFAULT;
        return -1;
    }

    char *filename = kmalloc(PATH_MAX);
    struct vnode* v;

    int instr_return = copyinstr(program, filename, PATH_MAX, NULL);
    if (instr_return) {
        *err = EFAULT;
        kfree(filename);
        return -1;
    }

    instr_return = vfs_open(filename, O_RDONLY, 0, &v);
    if (instr_return) {
        *err = instr_return;
        kfree(filename);
        return -1;
    }

    // Start with 4K size
    int arg_size = PAGE_SIZE;

    size_t data_length = 0;
    size_t arg_len = 0;
    int argCount = 0;
    userptr_t cur_arg;
    userptr_t shifted_args = args;
    char *argv;
ALLOC_ARGS:

    data_length = 0;
    arg_len = 0;
    argCount = 0;
    shifted_args = args;
    argv = kmalloc(arg_size);

    if(argv==NULL)
    {
        *err = ENOMEM;
        kfree(filename);
        return 1;
    }
    
    while(true)
    {
        // Get address of an argument
        instr_return = copyin(shifted_args, &cur_arg, sizeof(userptr_t));
        if (instr_return) {
            *err = EFAULT;
            vfs_close(v);
            kfree(argv);
            kfree(filename);
            return -1;
        }

        // if argument null, reached end of arguments
        if (cur_arg == NULL) {
			break;
		}
    
        // Copy argument contents into argument list argv
        instr_return = copyinstr(cur_arg, argv + data_length, arg_size - data_length, &arg_len);
        if(instr_return == ENAMETOOLONG)
        {
            // If 4K was not enough, try again with 64K
            kfree(argv);
            arg_size = ARG_MAX;
            goto ALLOC_ARGS;
        }
        else if (instr_return) {
            *err = EFAULT;
            vfs_close(v);
            kfree(argv);
            kfree(filename);
            return - 1;
        }

        // Update argument count and apply necessary shifts for next read
        argCount++;
        data_length += arg_len;
        shifted_args += sizeof(userptr_t);
    }

    vaddr_t entryAddr,stackAddr;
    struct addrspace *vm_backup;
    struct addrspace *vm_new;

    // Create a new address space for program
    vm_new = as_create();
    if (vm_new == NULL) {
        vfs_close(v);
        *err = ENOMEM;
        kfree(argv);
        kfree(filename);
        return -1;
    }

    /*
     * Create a backup of original address space to restore if needed 
     * and set current proc's address space to newly created one
     */
    spinlock_acquire(&curproc->p_lock);
    vm_backup = curproc->p_addrspace;
    curproc->p_addrspace = vm_new;
    spinlock_release(&curproc->p_lock);
    
    as_activate();

    // Load elf pointed to by program pointer
    instr_return = load_elf(v, &entryAddr);
    if (instr_return) {
        // Restore original address space and cleanup before exiting
        spinlock_acquire(&curproc->p_lock);
        curproc->p_addrspace = vm_backup;
        spinlock_release(&curproc->p_lock);
        vfs_close(v);     
        as_destroy(vm_new);        
        *err = instr_return;
        kfree(filename);
        return -1;
    }

    vfs_close(v);

    instr_return = as_define_stack(vm_new, &stackAddr);
    if (instr_return) {
        // Restore original address space and cleanup before exiting
        spinlock_acquire(&curproc->p_lock);
        curproc->p_addrspace = vm_backup;
        spinlock_release(&curproc->p_lock);
        as_destroy(vm_new);        
        *err = instr_return;
        kfree(filename);
        return -1;
    }

    /*
    * stackAddr initialized by DUMBVM to 0x80000000, marking the top of the kuseg user segment
    * Above the user stack, there will be a data segment and an argument segment containing arg data and pointers
    * Want to shift stackAddr down below these segments and populate these areas of the stack before passing to usermode
    * 0x80000000
    * 0x7fffffff┌────────────┐
    *           │            │ kuseg
    *           │            │
    *           │Data Segment│
    *           │            │
    *           │            │
    *           │            │
    *           │            │
    *           ├────────────┤
    *           │Arg Segment │
    *           │            │
    *           │            │
    *           │            │
    *           │            │
    *           │            │
    *  stackAddr├────────────┤
    *           │            │
    *           │            │
    *           │            │
    *           │            │
    *           └────────────┘
    */

    // Use stack pointer as starting address to determine address locations of data and args
    // Adjust stack pointer depending on data size and mark bottom of data segment, data segment stored above args
    stackAddr -= data_length;
    // Ensure that stack address is aligned
    stackAddr -= (stackAddr & (sizeof(void *) - 1));
    userptr_t user_stack_data_vaddr = (userptr_t)stackAddr;

    // Move stack down to make space for args and mark bottom of arg segment
    stackAddr -=(argCount + 1) * sizeof(userptr_t);
    userptr_t user_stack_args_vaddr = (userptr_t)stackAddr;

    for(size_t offset = 0; offset < data_length; offset += arg_len)
    {
        cur_arg =  user_stack_data_vaddr + offset;

        // copy out adress to argument to userspace
        instr_return = copyout(&cur_arg, user_stack_args_vaddr, sizeof(cur_arg));
        if (instr_return)
        {
            spinlock_acquire(&curproc->p_lock);
            curproc->p_addrspace = vm_backup;
            spinlock_release(&curproc->p_lock);
            as_destroy(vm_new);        
            *err = instr_return;
            kfree(filename);
            return -1;
        }

        // copy out argument to userspace
        instr_return = copyoutstr(argv + offset, cur_arg, data_length - offset , &arg_len);
        if (instr_return)
        {
            spinlock_acquire(&curproc->p_lock);
            curproc->p_addrspace = vm_backup;
            spinlock_release(&curproc->p_lock);
            as_destroy(vm_new);        
            *err = instr_return;
            kfree(filename);
            return -1;
        }
        // Advance args vaddr pointer by size of written argument
        user_stack_args_vaddr += sizeof(cur_arg);
    }

    // Add null entry to end of arg list
    cur_arg = NULL;
    instr_return = copyout(&cur_arg, user_stack_args_vaddr, sizeof(cur_arg));
    if (instr_return)
    {
        spinlock_acquire(&curproc->p_lock);
        curproc->p_addrspace = vm_backup;
        spinlock_release(&curproc->p_lock);
        as_destroy(vm_new);        
        *err = instr_return;
        kfree(filename);
        return -1;
    }

    // Backup no longer needed, cannot allow execv to fail after this
    if(vm_backup)
    {
        as_destroy(vm_backup);
    }

    // Update current thread's name based on input
    kfree(curthread->t_name);
    curthread->t_name = kstrdup(filename);
    kfree(filename);
	
	enter_new_process(argCount, (userptr_t)stackAddr, NULL, stackAddr, entryAddr);

    panic("returned from new process");
	
	*err = EINVAL;
	return -1;
}