#include "types.h"
#include "traps.h"
#include "param.h"
#include "memlayout.h"
#include "mmu.h"
#include "proc.h"
#include "defs.h"
#include "x86.h"
#include "elf.h"
#include "fcntl.h"
#include "errno.h"
#include "x86.h"
#include "fsstat.h"

// Binaries with entry points at 0x100000+ are musl/Linux-ABI.
// SchOS-native binaries (init) are linked at 0x4000.
#define LINUX_ABI_ENTRY_THRESHOLD 0x100000UL
extern int linux_syscall_map[];
extern int linux_syscall_map_size;

// envp may be null, which means "no environment" -- the same thing this
// function used to do unconditionally.
int
exec(char *path, char **argv, char **envp)
{
  uint64 _t0 = rdtsc();
  fsstat.exec_calls++;
  char *s, *last;
  int i, off;
  uint64 argc, envc, sz, sp;
  struct elfhdr elf;
  struct inode *ip;
  struct proghdr ph;
  pml4e_t *pgdir, *oldpgdir;
  int slot;

  // The three pointer arrays live on the heap, not on the kernel stack.
  //
  // Together they are about 1.8 KB, and KSTACKSIZE is 4096 -- with the call
  // chain under this function running exec -> loaduvm -> readi -> SchFS ->
  // btree, there is not 1.8 KB to spare. Putting them here overflowed the
  // stack on the first exec after the environment was added, and the canary
  // in param.h caught it on pid 1 before it could corrupt anything.
  //
  // One page holds all three with room over.
  uint64 *scratch = (uint64*)kalloc();
  if(!scratch)
    return -1;
  uint64 *ustack      = scratch;                      /* MAXARG+1  */
  uint64 *estack      = ustack + MAXARG + 1;          /* MAXENV+1  */
  uint64 *stack_slots = estack + MAXENV + 1;          /* the rest   */
  _Static_assert((MAXARG + 1 + MAXENV + 1 + MAXARG + MAXENV + 32) * 8 <= 4096,
                 "exec scratch does not fit in one page");
  uint64 max_sz = 0;
  uint64 rand_va = 0;

  begin_op();

  if((ip = namei(path)) == 0){
    end_op();
    // Must be ENOENT, not a bare -1: execvp() relies on this exact errno to
    // know it should try the next $PATH component rather than give up.
    return -ENOENT;
  }

  ilock(ip);
  pgdir = 0;

  if(readi(ip, 0, (uint64)&elf, 0, sizeof(elf)) != sizeof(elf))
    goto bad;

  if(elf.magic != ELF_MAGIC)
    goto bad;

  if((pgdir = setupkvm()) == 0)
    goto bad;

  for(i=0, off=elf.phoff; i<elf.phnum; i++, off+=elf.phentsize){
    if(readi(ip, 0, (uint64)&ph, off, sizeof(ph)) != sizeof(ph))
      goto bad;
    if(ph.type != ELF_PROG_LOAD)
      continue;
    if(ph.memsz < ph.filesz)
      goto bad;
    if(ph.vaddr + ph.memsz < ph.vaddr)
      goto bad;

    uint64 start = ph.vaddr & ~0xFFF;
    uint64 end = ph.vaddr + ph.memsz;
    if(allocuvm(pgdir, start, end) == 0)
      goto bad;

    if(loaduvm(pgdir, (char*)ph.vaddr, ip, ph.off, ph.filesz) < 0)
      goto bad;

    if(end > max_sz)
      max_sz = end;
  }
  iunlockput(ip);
  end_op();
  ip = 0;

  sz = PGROUNDUP(max_sz);
  if((sz = allocuvm(pgdir, sz, sz + 64*PGSIZE)) == 0)
    goto bad;
  clearpteu(pgdir, (char*)(sz - 64*PGSIZE));
  sp = sz;

  for(argc = 0; argv[argc]; argc++) {
    if(argc >= MAXARG)
      goto bad;
    sp -= strlen(argv[argc]) + 1;
    sp &= ~7;
    if(copyout(pgdir, sp, argv[argc], strlen(argv[argc]) + 1) < 0)
      goto bad;
    ustack[argc] = sp;
  }

  // The environment, copied the same way and for the same reason: these
  // strings live in the *old* address space, which is still mapped, and will
  // not be once the new pgdir is installed.
  for(envc = 0; envp && envp[envc]; envc++) {
    if(envc >= MAXENV)
      goto bad;
    sp -= strlen(envp[envc]) + 1;
    sp &= ~7;
    if(copyout(pgdir, sp, envp[envc], strlen(envp[envc]) + 1) < 0)
      goto bad;
    estack[envc] = sp;
  }

  // Push 16 pseudo-random bytes for AT_RANDOM.
  {
    uint8 randbuf[16];
    uint rseed = (uint)(uintptr_t)myproc() ^ (uint)myproc()->pid;
    for(int ri = 0; ri < 16; ri++){
      rseed = rseed * 1664525u + 1013904223u;
      randbuf[ri] = (uint8)(rseed >> 16);
    }
    sp -= 16; sp &= ~0xFUL;
    rand_va = sp;
    if(copyout(pgdir, rand_va, (char*)randbuf, 16) < 0) goto bad;
  }

  slot = 0;
  stack_slots[slot++] = argc;
  for(i = 0; i < argc; i++) {
    stack_slots[slot++] = ustack[i];
  }
  stack_slots[slot++] = 0;    // end of argv

  // The environment used to be a single zero here, with the comment "(no
  // environment)". Nothing on SchOS had one: not a shell, not a musl program
  // that had just called execve with a perfectly good envp, not anything. So
  // $TERM was unset, $PATH was whatever each program invented, and a variable
  // exported in a shell did not survive the exec that was supposed to carry it.
  for(i = 0; i < envc; i++)
    stack_slots[slot++] = estack[i];
  stack_slots[slot++] = 0;    // end of envp

  // Auxiliary vector — musl reads these for page size, clock, etc.
  stack_slots[slot++] = 6;    // AT_PAGESZ
  stack_slots[slot++] = 4096;
  stack_slots[slot++] = 17;   // AT_CLKTCK
  stack_slots[slot++] = USER_HZ;   // what times() is scaled to, not a literal
  stack_slots[slot++] = 23;   // AT_SECURE = 0 (not setuid)
  stack_slots[slot++] = 0;
  stack_slots[slot++] = 11;   // AT_UID
  stack_slots[slot++] = 0;
  stack_slots[slot++] = 12;   // AT_EUID
  stack_slots[slot++] = 0;
  stack_slots[slot++] = 13;   // AT_GID
  stack_slots[slot++] = 0;
  stack_slots[slot++] = 14;   // AT_EGID
  stack_slots[slot++] = 0;
  stack_slots[slot++] = 16;   // AT_HWCAP
  stack_slots[slot++] = 0;
  stack_slots[slot++] = 25;   // AT_RANDOM
  stack_slots[slot++] = rand_va;
  stack_slots[slot++] = 0;    // AT_NULL
  stack_slots[slot++] = 0;

  sp -= slot * sizeof(uint64);
  sp &= ~0xF;

  if(copyout(pgdir, sp, (char*)stack_slots, slot * sizeof(uint64)) < 0)
    goto bad;

  for(last=s=path; *s; s++)
    if(*s == '/')
      last = s+1;
  safestrcpy(myproc()->name, last, sizeof(myproc()->name));
  safestrcpy(myproc()->exe, path, sizeof(myproc()->exe));

  for(int i = 0; i < NOFILE; i++){
    if(myproc()->ofile[i] && (myproc()->ofile_flags[i] & FD_CLOEXEC)){
      fileclose(myproc()->ofile[i]);
      myproc()->ofile[i] = 0;
      myproc()->ofile_flags[i] = 0;
    }
  }

  // Auto-select syscall ABI: musl binaries have entry >= 0x100000;
  // SchOS-native binaries (init) are linked at 0x4000.
  if (elf.entry >= LINUX_ABI_ENTRY_THRESHOLD) {
    myproc()->syscall_map = linux_syscall_map;
    myproc()->map_size    = linux_syscall_map_size;
  } else {
    myproc()->syscall_map = 0;
    myproc()->map_size    = 0;
  }

  // Detach shared segments before anything forgets which mappings they were.
  //
  // This has to happen while myproc()->pgdir and myproc()->vma[] still describe the
  // *old* image: the memset below erases the bookkeeping, and freevm() at the
  // end would then free a shared memory segment's pages as though they were
  // ordinary heap. Doing it here rather than there is the difference between
  // detaching a segment and corrupting whoever else has it mapped.
  proc_release_shared(myproc());

  oldpgdir = myproc()->pgdir;
  myproc()->pgdir = pgdir;
  myproc()->sz = sz;

  // New address space: reset all per-process mmap/signal/TLS state.
  // sig_handler/sig_restorer point into the *old* image — POSIX resets
  // handled signals to default on exec, and keeping them would make the
  // first delivered signal jump to a stale address in the new image.
  // mmap_next likewise belongs to the old address space; left alone it
  // could hand out mmap ranges overlapping the new image's layout.
  memset(myproc()->vma, 0, sizeof(myproc()->vma));
  memset(myproc()->sig_handler, 0, sizeof(myproc()->sig_handler));
  memset(myproc()->sig_restorer, 0, sizeof(myproc()->sig_restorer));
  myproc()->fs_base          = 0;
  myproc()->signal_mask      = 0;
  myproc()->pending_signals  = 0;
  myproc()->in_signal        = 0;
  myproc()->sigsusp_active   = 0;
  myproc()->mmap_next        = 0;

  myproc()->tf->rax = 0;
  myproc()->tf->rbx = 0;
  myproc()->tf->rdx = 0;
  myproc()->tf->rbp = 0;
  myproc()->tf->r8  = 0;
  myproc()->tf->r9  = 0;
  myproc()->tf->r10 = 0;
  myproc()->tf->r11 = 0x200; // RFLAGS for sysretq: IF=1 (bit 9) so user gets timer IRQs
  myproc()->tf->r12 = 0;
  myproc()->tf->r13 = 0;
  myproc()->tf->r14 = 0;
  myproc()->tf->r15 = 0;

  myproc()->tf->rdi = argc;
  myproc()->tf->rsi = sp + 8;
  myproc()->tf->rip = elf.entry;
  myproc()->tf->rcx = elf.entry;
  myproc()->tf->rsp = sp;

  // syscall_entry saves user RSP at kstack+PGSIZE-16 (not in the trapframe).
  // sysretq reads RSP from that slot, so update it to the new stack pointer.
  // Without this, sysretq returns with the old parent's RSP, causing lksh to
  // start with an invalid/wrong stack pointer.
  *(uint64*)((char*)myproc()->kstack + PGSIZE - 16) = sp;

  // Userspace is running now, so the splash has done its job and the display
  // becomes a terminal. Only the first exec does anything; every later one
  // finds the console already in charge.
  console_takeover();

  switchuvm(myproc());
  freevm(oldpgdir);
  kfree((char*)scratch);
  fsstat.exec_cycles += rdtsc() - _t0;
  return 0;

 bad:
  kfree((char*)scratch);
  if(pgdir)
    freevm(pgdir);
  if(ip){
    iunlockput(ip);
    end_op();
  }
  return -1;
}