#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"
#include "spinlock.h"
#include "sleeplock.h"
#include "fs.h"
#include "file.h"
#include "uaccess.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;

// Where the argument strings come from: an array in the kernel (the first
// exec, of /bin/schinit), or one in the calling process's memory.
//
// The user case reads every pointer and every string through the uaccess
// helpers, and copies through a kernel buffer, rather than being handed
// pointers somebody checked earlier. Another thread of the same process can
// rewrite the array between a check and a use, and a pointer swapped for a
// kernel address afterwards would have copied kernel memory into the new
// program's stack.
struct argsrc {
  int    user;
  char **k;
  addr_t u;
};

// Pointer i of the array, 0 past its end; -1 on a fault.
static int
arg_ptr(struct argsrc *s, uint64 i, addr_t *p)
{
  if (s->user)
    return s->u ? fetchaddr(s->u + i * sizeof(addr_t), p) : (*p = 0, 0);
  *p = s->k ? (addr_t)s->k[i] : 0;
  return 0;
}

static long
arg_len(struct argsrc *s, addr_t p)
{
  return s->user ? strlen_user(p) : (long)strlen((char *)p);
}

// Count an array and its bytes, against Linux's limits: ARG_MAX in all,
// pointers included, and ARG_STRLEN_MAX for any one string.
static int
arg_measure(struct argsrc *s, uint64 *count, uint64 *bytes, uint64 already)
{
  for (uint64 i = 0; ; i++) {
    addr_t p;
    if (arg_ptr(s, i, &p) < 0)
      return -EFAULT;
    if (p == 0) {
      *count = i;
      return 0;
    }
    long n = arg_len(s, p);
    if (n < 0)
      return n == -EFAULT ? -EFAULT : -E2BIG;
    if (n + 1 > ARG_STRLEN_MAX)
      return -E2BIG;
    *bytes += (uint64)n + 1;
    if (*bytes + (already + i + 1) * sizeof(addr_t) > ARG_MAX)
      return -E2BIG;
  }
}

// The new stack's pointer table, written in order a page at a time.
struct tbl {
  pml4e_t *pgdir;
  uint64  *buf;
  addr_t   at;          // where buf[0] goes in the new address space
  int      n;
};

static int
tbl_flush(struct tbl *t)
{
  if (t->n && copyout(t->pgdir, t->at, (char *)t->buf, t->n * 8) < 0)
    return -1;
  t->at += (addr_t)t->n * 8;
  t->n = 0;
  return 0;
}

static int
tbl_put(struct tbl *t, uint64 v)
{
  t->buf[t->n++] = v;
  return t->n == PGSIZE / 8 ? tbl_flush(t) : 0;
}

// Copy `count` strings to `*at` onward in the new address space, each NUL
// terminated whatever happened to it since it was measured, and put each
// one's address in the table. The strings must not have grown: the room for
// them was measured once, and running past it would write over the table.
static int
arg_copy(struct argsrc *s, uint64 count, pml4e_t *pgdir, addr_t *at,
         addr_t end, char *bounce, struct tbl *t)
{
  for (uint64 i = 0; i < count; i++) {
    addr_t p;
    if (arg_ptr(s, i, &p) < 0 || p == 0)
      return -1;
    long n = arg_len(s, p);
    if (n < 0 || *at + (addr_t)n + 1 > end)
      return -1;
    if (tbl_put(t, *at) < 0)
      return -1;
    for (long done = 0; done < n; ) {
      long c = n - done > PGSIZE ? PGSIZE : n - done;
      if (s->user) {
        if (copy_from_user(bounce, p + done, c) < 0)
          return -1;
      } else {
        memmove(bounce, (char *)p + done, c);
      }
      if (copyout(pgdir, *at + done, bounce, c) < 0)
        return -1;
      done += c;
    }
    char nul = 0;
    if (copyout(pgdir, *at + n, &nul, 1) < 0)
      return -1;
    *at += (addr_t)n + 1;
  }
  return 0;
}

static int do_exec(char *path, struct argsrc *av, struct argsrc *ev);

// 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)
{
  struct argsrc a = { 0, argv, 0 }, e = { 0, envp, 0 };
  return do_exec(path, &a, &e);
}

// execve(2): the arrays are in the caller's memory. See struct argsrc.
int
exec_user(char *path, addr_t uargv, addr_t uenvp)
{
  struct argsrc a = { 1, 0, uargv }, e = { 1, 0, uenvp };
  return do_exec(path, &a, &e);
}

static int
do_exec(char *path, struct argsrc *av, struct argsrc *ev)
{
  uint64 _t0 = rdtsc();
  fsstat.exec_calls++;
  char *s, *last;
  int i, off;
  uint64 sz, sp;
  struct elfhdr elf;
  struct inode *ip;
  struct proghdr ph;
  pml4e_t *pgdir, *oldpgdir;

  // The arguments are measured before anything is loaded, so that too many
  // of them is E2BIG with nothing to undo -- and so the stack can be sized to
  // hold them.
  //
  // There used to be a fixed 32 arguments and 64 variables, in arrays on the
  // kernel stack and then in one page. `rm *.o` in a build directory, or any
  // glob over more than 31 files, failed with E2BIG. Now the limit is Linux's:
  // 2 MB of strings and pointers together.
  uint64 argc = 0, envc = 0, strbytes = 0;
  int r;
  if ((r = arg_measure(av, &argc, &strbytes, 0)) < 0 ||
      (r = arg_measure(ev, &envc, &strbytes, argc + 1)) < 0)
    return r;

  // Two pages, off the kernel stack, which has no room for either: a bounce
  // buffer for the strings and a batch of pointer-table entries.
  char *bounce = kalloc();
  uint64 *scratch = (uint64*)kalloc();
  if(!bounce || !scratch){
    if (bounce) kfree(bounce);
    if (scratch) kfree((char*)scratch);
    return -ENOMEM;
  }
  uint64 max_sz = 0;
  uint64 rand_va = 0;

  begin_op();

  int err;
  if((ip = nameifollow(0, path, &err)) == 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 -err;
  }

  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;

  // The stack: 63 pages plus a guard, as always, and on top of that room for
  // the arguments, the pointer table and AT_RANDOM.
  uint64 nslots = 1 + argc + 1 + envc + 1 + 20;
  uint64 argpages = (strbytes + nslots * 8 + 64 + PGSIZE - 1) / PGSIZE;
  uint64 stackpages = 64 + argpages;
  sz = PGROUNDUP(max_sz);
  if((sz = allocuvm(pgdir, sz, sz + stackpages*PGSIZE)) == 0)
    goto bad;
  clearpteu(pgdir, (char*)(sz - stackpages*PGSIZE));

  // From the top: the strings, then 16 random bytes, then the table that
  // points at them -- argc, argv, 0, envp, 0, the auxiliary vector -- with
  // the stack pointer at argc, 16-byte aligned, as the ELF ABI requires.
  addr_t strbase = (sz - strbytes) & ~0xFULL;
  rand_va = (strbase - 16) & ~0xFULL;
  sp = (rand_va - nslots * 8) & ~0xFULL;

  {
    uint8 randbuf[16];
    uint rseed = (uint)(uintptr_t)myproc() ^ (uint)myproc()->pid ^ (uint)rdtsc();
    for(int ri = 0; ri < 16; ri++){
      rseed = rseed * 1664525u + 1013904223u;
      randbuf[ri] = (uint8)(rseed >> 16);
    }
    if(copyout(pgdir, rand_va, (char*)randbuf, 16) < 0) goto bad;
  }

  struct tbl t = { pgdir, scratch, sp, 0 };
  addr_t at = strbase;
  if (tbl_put(&t, argc) < 0)
    goto bad;
  if (arg_copy(av, argc, pgdir, &at, sz, bounce, &t) < 0 || tbl_put(&t, 0) < 0)
    goto bad;
  // 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.
  if (arg_copy(ev, envc, pgdir, &at, sz, bounce, &t) < 0 || tbl_put(&t, 0) < 0)
    goto bad;

  // Auxiliary vector — musl reads these for page size, clock, etc.
  {
    static const uint64 fixed[] = {
      6, 4096,          // AT_PAGESZ
      17, USER_HZ,      // AT_CLKTCK: what times() is scaled to, not a literal
      23, 0,            // AT_SECURE: not setuid
      11, 0, 12, 0,     // AT_UID, AT_EUID
      13, 0, 14, 0,     // AT_GID, AT_EGID
      16, 0,            // AT_HWCAP
    };
    for (int k = 0; k < (int)NELEM(fixed); k++)
      if (tbl_put(&t, fixed[k]) < 0)
        goto bad;
    if (tbl_put(&t, 25) < 0 || tbl_put(&t, rand_va) < 0 ||   // AT_RANDOM
        tbl_put(&t, 0) < 0 || tbl_put(&t, 0) < 0 ||          // AT_NULL
        tbl_flush(&t) < 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);
  kfree(bounce);
  fsstat.exec_cycles += rdtsc() - _t0;
  return 0;

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