// TODO: CHECK/REPLACE/UPDATE OLD CODE (this file is based on xv6)
#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "riscv.h"
#include "sched.h"
#include "proc.h"
#include "defs.h"
#include "drives.h"
#include "fsinstance.h"
//#include "elf.h"
#include <libc/elf.h>
#include "kprintf.h"

static int loadseg(pde_t *, uint64, fsinstance_inode_t *, uint, uint);

int flags2perm(int flags)
{
    int perm = 0;
    if(flags & 0x1)
      perm = PTE_X;
    if(flags & 0x2)
      perm |= PTE_W;
    return perm;
}

// TODO: This should probably also look up programs from PATH or else
// be execv instead, check manpages or standards thoroughly for correct
// behaviour/names.
int
execve(char *path, char **argv, char **envv)
{
  //printf("Attempting exec of '%s'\n", path);
  char *s, *last;
  int i, off;
  uint64 argc, envc, sz = 0, sp, ustack[MAXARG*2+2], stackbase;
  Elf64_Ehdr elf;
  fsinstance_inode_t *ip;
  Elf64_Phdr ph;
  pagetable_t pagetable = 0, oldpagetable;
  struct proc *p = myproc();

  fsinstance_t* instance = drives_fsbegin(p->drives, p->cwdrive, path);

  int typevar;
  if((ip = drives_fsnode(p->drives, p->cwdrive, path, &typevar)) == NULL || typevar != DRIVES_HANDLER_FS){
    printf("BAD EXEC\n");
    drives_fsend(p->drives, instance);
    return -1;
  }
  fsinstance_inode_lockandload(ip);

  // Check ELF header
  if(fsinstance_inode_read(ip, 0, (uint64)&elf, 0, sizeof(elf)) != sizeof(elf))
    goto bad;

  if (elf.e_ident[0] != ELFMAG0 || elf.e_ident[1] != ELFMAG1 || elf.e_ident[2] != ELFMAG2 || elf.e_ident[3] != ELFMAG3) {
    goto bad;
  }

  if((pagetable = proc_pagetable(p)) == 0)
    goto bad;

  // Load program into memory.
  for(i=0, off=elf.e_phoff; i<elf.e_phnum; i++, off+=sizeof(ph)){
    if(fsinstance_inode_read(ip, 0, (uint64)&ph, off, sizeof(ph)) != sizeof(ph))
      goto bad;
    if(ph.p_type != PT_LOAD) {
      continue;
    }
    if(ph.p_memsz < ph.p_filesz)
      goto bad;
    if(ph.p_vaddr + ph.p_memsz < ph.p_vaddr)
      goto bad;
    if(ph.p_vaddr % PGSIZE != 0)
      goto bad;
    uint64 sz1;
    if((sz1 = uvmalloc(pagetable, sz, ph.p_vaddr + ph.p_memsz, flags2perm(ph.p_flags))) == 0)
      goto bad;
    sz = sz1;
    if(loadseg(pagetable, ph.p_vaddr, ip, ph.p_offset, ph.p_filesz) < 0)
      goto bad;
  }
  fsinstance_inode_unlockandunget(ip);
  drives_fsend(p->drives, instance);
  ip = 0;

  p = myproc();
  uint64 oldsz = p->sz;

  // Allocate some pages at the next page boundary.
  // Make the first inaccessible as a stack guard.
  // Use the rest as the user stack.
  sz = PGROUNDUP(sz);
  uint64 sz1;
  if((sz1 = uvmalloc(pagetable, sz, sz + (USERSTACK+1)*PGSIZE, PTE_W)) == 0)
    goto bad;
  sz = sz1;
  uvmclear(pagetable, sz-(USERSTACK+1)*PGSIZE);
  sp = sz;
  stackbase = sp - USERSTACK*PGSIZE;

  // Push argument strings, prepare rest of stack in ustack.
  for(argc = 0; argv[argc]; argc++) {
    if(argc >= MAXARG)
      goto bad;
    sp -= strlen(argv[argc]) + 1;
    sp -= sp % 16; // riscv sp must be 16-byte aligned
    if(sp < stackbase)
      goto bad;
    if(copyout(pagetable, sp, argv[argc], strlen(argv[argc]) + 1) < 0)
      goto bad;
    ustack[argc] = sp;
  }
  ustack[argc] = 0;
  // Push env strings
  for(envc = 0; envv[envc]; envc++) {
    if(envc >= MAXARG)
      goto bad;
    sp -= strlen(envv[envc]) + 1;
    sp -= sp % 16; // riscv sp must be 16-byte aligned
    if(sp < stackbase)
      goto bad;
    if(copyout(pagetable, sp, envv[envc], strlen(envv[envc]) + 1) < 0)
      goto bad;
    ustack[argc+envc] = sp;
  }
  ustack[argc+envc] = 0;

  // push the array of argv[] pointers.
  sp -= (argc+1+envc+1) * sizeof(uint64);
  sp -= sp % 16;
  if(sp < stackbase)
    goto bad;
  if(copyout(pagetable, sp, (char *)ustack, (argc+1+envc+1)*sizeof(uint64)) < 0)
    goto bad;

  // arguments to user main(argc, argv)
  // argc is returned via the system call return
  // value, which goes in a0.
  p->trapframe->a1 = sp;
  p->trapframe->a2 = sp + (argc+1)*sizeof(uint64);

  // Save program name for debugging.
  for(last=s=path; *s; s++)
    if(*s == '/')
      last = s+1;
  safestrcpy(p->name, last, PROC_NAME_SIZE /*sizeof(p->name) TODO */);
    
  // Commit to the user image.
  oldpagetable = p->pagetable;
  p->pagetable = pagetable;
  p->sz = sz;
  p->trapframe->epc = elf.e_entry;  // initial program counter = main
  p->trapframe->sp = sp; // initial stack pointer
  proc_freepagetable(oldpagetable, oldsz, p->mainthread == 0);
  p->mainthread = 0ULL;

  return argc; // this ends up in a0, the first argument to main(argc, argv)

 bad:
 printf("BAD EXEC\n");
  if(pagetable)
    proc_freepagetable(pagetable, sz, 1);
  if(ip){
    fsinstance_inode_unlockandunget(ip);
    drives_fsend(p->drives, instance);
  }
  return -1;
}

// Load a program segment into pagetable at virtual address va.
// va must be page-aligned
// and the pages from va to va+sz must already be mapped.
// Returns 0 on success, -1 on failure.
static int
loadseg(pagetable_t pagetable, uint64 va, fsinstance_inode_t *ip, uint offset, uint sz)
{
  uint i, n;
  uint64 pa;

  for(i = 0; i < sz; i += PGSIZE){
    pa = walkaddr(pagetable, va + i);
    if(pa == 0)
      panic("loadseg: address should exist");
    if(sz - i < PGSIZE)
      n = sz - i;
    else
      n = PGSIZE;
    if(fsinstance_inode_read(ip, 0, (uint64)pa, offset+i, n) != n)
      return -1;
  }
  
  return 0;
}