#include "core/types.h"
#include "arch/traps.h"
#include "core/param.h"
#include "arch/memlayout.h"
#include "arch/mmu.h"
#include "core/proc.h"
#include "core/defs.h"
#include "arch/x86.h"
#include "core/elf.h"
#include "uapi/fcntl.h"
#include "uapi/errno.h"
#include "arch/x86.h"
#include "fs/fsstat.h"
#include "core/spinlock.h"
#include "core/sleeplock.h"
#include "fs/fs.h"
#include "fs/file.h"
#include "core/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.
//
// `pre` goes in front, always from the kernel, and `skip` drops that many of
// the array's own: a #! script runs as `interp [arg] script argv[1]...`, which
// is the caller's array less its argv[0] behind three strings of ours.
struct argsrc {
  int    user;
  char **k;
  addr_t u;
  char  *pre[3];
  int    npre;
  uint64 skip;
};

// Pointer i of the array, 0 past its end; -1 on a fault. *kern says which
// side of the boundary it points to.
static int
arg_ptr(struct argsrc *s, uint64 i, addr_t *p, int *kern)
{
  if (i < (uint64)s->npre) {
    *p = (addr_t)s->pre[i];
    *kern = 1;
    return 0;
  }
  i = i - s->npre + s->skip;
  *kern = !s->user;
  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(addr_t p, int kern)
{
  return kern ? (long)strlen((char *)p) : strlen_user(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;
    int kern;
    if (arg_ptr(s, i, &p, &kern) < 0)
      return -EFAULT;
    if (p == 0) {
      *count = i;
      return 0;
    }
    long n = arg_len(p, kern);
    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;
    int kern;
    if (arg_ptr(s, i, &p, &kern) < 0 || p == 0)
      return -1;
    long n = arg_len(p, kern);
    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 (!kern) {
        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);
static int exec_elf(char *path, struct argsrc *av, struct argsrc *ev);

// --- Dynamic linking ---------------------------------------------------------
//
// A dynamically linked program names an interpreter in PT_INTERP -- on musl,
// /lib/ld-musl-x86_64.so.1, which is libc.so under another name -- and the
// kernel's whole part in dynamic linking is to load that as well and start it
// instead of the program. The interpreter finds the program through the
// auxiliary vector (AT_PHDR, AT_ENTRY), maps libc and anything else it needs
// with mmap(), relocates, and jumps to AT_ENTRY itself.
//
// Position-independent images -- the interpreter always, a PIE program
// sometimes -- are ET_DYN, linked at zero and placed at a base of the
// kernel's choosing. A PIE program goes where a fixed one would; the
// interpreter goes at the start of the mmap area, and mmap carries on above
// it, so the two never meet.
#define ET_DYN       3
#define PT_INTERP    3
#define PT_PHDR      6
#define PIE_BASE     0x400000UL      // where lld puts a fixed executable
#define INTERP_BASE  0x40000000UL    // the bottom of sys_mmap's bump area

// Map and read one image's PT_LOAD segments at `base`. Returns 0 and the end
// of the highest segment, and where the program headers are once loaded (0 if
// no segment carries them), or -1.
//
// Segments are mapped from where the previous one's last page ended: two
// segments may share a page, and asking allocuvm() for it twice would map it
// twice.
static int
load_image(pml4e_t *pgdir, struct inode *ip, struct elfhdr *elf, uint64 base,
           uint64 *hi, uint64 *phdr_va)
{
  struct proghdr ph;
  uint64 mapped = 0;

  *hi = 0;
  *phdr_va = 0;
  for (int i = 0, off = elf->phoff; i < elf->phnum; i++, off += elf->phentsize) {
    if (readi(ip, 0, (uint64)&ph, off, sizeof(ph)) != sizeof(ph))
      return -1;
    if (ph.type == PT_PHDR)
      *phdr_va = base + ph.vaddr;
    if (ph.type != ELF_PROG_LOAD)
      continue;
    if (ph.memsz < ph.filesz || ph.vaddr + ph.memsz < ph.vaddr)
      return -1;

    uint64 start = (base + ph.vaddr) & ~0xFFFULL;
    uint64 end = base + ph.vaddr + ph.memsz;
    if (start < mapped)
      start = mapped;
    if (start < end && allocuvm(pgdir, start, end) == 0)
      return -1;
    if (PGROUNDUP(end) > mapped)
      mapped = PGROUNDUP(end);

    if (loaduvm(pgdir, (char*)(base + ph.vaddr), ip, ph.off, ph.filesz) < 0)
      return -1;

    // The headers are in the file at phoff; whichever segment covers that
    // offset is where they are in memory. PT_PHDR says so directly when the
    // linker wrote one, and this covers the images that do not.
    if (!*phdr_va && ph.off <= elf->phoff && elf->phoff < ph.off + ph.filesz)
      *phdr_va = base + ph.vaddr + (elf->phoff - ph.off);

    if (end > *hi)
      *hi = end;
  }
  return 0;
}

// The interpreter path from PT_INTERP, or "" if there is none. -1 on a
// malformed one.
static int
find_interp(struct inode *ip, struct elfhdr *elf, char *buf, int cap)
{
  struct proghdr ph;

  buf[0] = 0;
  for (int i = 0, off = elf->phoff; i < elf->phnum; i++, off += elf->phentsize) {
    if (readi(ip, 0, (uint64)&ph, off, sizeof(ph)) != sizeof(ph))
      return -1;
    if (ph.type != PT_INTERP)
      continue;
    if (ph.filesz < 2 || ph.filesz > (uint64)cap)
      return -1;
    if (readi(ip, 0, (uint64)buf, ph.off, ph.filesz) != (int)ph.filesz)
      return -1;
    buf[ph.filesz - 1] = 0;
    return 0;
  }
  return 0;
}

// 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);
}

// --- #! scripts --------------------------------------------------------------
//
// A file that starts with #! names the program that runs it, and one optional
// argument, on its first line. Without this every script had to be started
// as `sh script`; apk runs its packages' triggers and install scripts with
// execve() directly, and each one failed with "Operation not permitted".
//
// If `path` is a script, puts "interp\0arg\0" in `line` (a page) and returns
// 1, with *argp pointing at the argument or 0. 0 if it is not one; <0 errno.
static int
script_interp(char *path, char *line, char **argp)
{
  struct inode *ip;
  int err, n;

  begin_op();
  if ((ip = nameifollow(0, path, &err)) == 0) {
    end_op();
    return -err;
  }
  ilock(ip);
  n = readi(ip, 0, (uint64)line, 0, 256);
  iunlockput(ip);
  end_op();
  if (n < 2 || line[0] != '#' || line[1] != '!')
    return 0;

  // Linux's rules: the line ends at a newline or after 255 bytes; blanks
  // separate the interpreter from the argument, and everything after the
  // first blank run is one argument, spaces and all.
  line[n < 256 ? n : 255] = 0;
  char *nl = line;
  while (*nl && *nl != '\n')
    nl++;
  *nl = 0;
  char *q = line + 2;
  while (*q == ' ' || *q == '\t')
    q++;
  if (*q == 0)
    return -ENOEXEC;
  char *ie = q;
  while (*ie && *ie != ' ' && *ie != '\t')
    ie++;
  char *a = ie;
  while (*a == ' ' || *a == '\t')
    a++;
  char *ae = a + strlen(a);
  while (ae > a && (ae[-1] == ' ' || ae[-1] == '\t' || ae[-1] == '\r'))
    ae--;
  *ae = 0;
  *ie = 0;
  memmove(line, q, ie - q + 1);
  if (*a) {
    int il = ie - q + 1;
    memmove(line + il, a, ae - a + 1);
    *argp = line + il;
  } else {
    *argp = 0;
  }
  return 1;
}

static int
do_exec(char *path, struct argsrc *av, struct argsrc *ev)
{
  char *page = kalloc();
  if (!page)
    return -ENOMEM;
  char *line = page, *spath = page + 2048, *arg;
  int r = script_interp(path, line, &arg);
  if (r <= 0) {
    kfree(page);
    return r < 0 ? r : exec_elf(path, av, ev);
  }

  // One level: an interpreter that is itself a script is ENOEXEC, which
  // exec_elf() says by itself when it finds no ELF header there.
  struct argsrc a = *av;
  safestrcpy(spath, path, 2048);
  a.npre = 0;
  a.pre[a.npre++] = line;
  if (arg)
    a.pre[a.npre++] = arg;
  a.pre[a.npre++] = spath;
  a.skip = av->skip + 1;
  r = exec_elf(line, &a, ev);
  kfree(page);
  return r;
}

static int
exec_elf(char *path, struct argsrc *av, struct argsrc *ev)
{
  uint64 _t0 = rdtsc();
  fsstat.exec_calls++;
  char *s, *last;
  uint64 sz, sp;
  struct elfhdr elf;
  struct inode *ip;
  pml4e_t *pgdir;
  struct mm *oldmm, *newmm;

  // 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;
  uint64 base = 0, phdr_va = 0, interp_end = 0, interp_entry = 0;
  char interp[128];

  begin_op();
  int in_op = 1;          // the error path closes the transaction only once
  int ret = -1;

  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)){
    ret = -ENOEXEC;           // too short to be anything we can run
    goto bad;
  }

  // Neither ELF nor, by now, a script. ENOEXEC is what tells musl's
  // execvp() to try it with /bin/sh.
  if(elf.magic != ELF_MAGIC){
    ret = -ENOEXEC;
    goto bad;
  }

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

  base = elf.type == ET_DYN ? PIE_BASE : 0;
  if (find_interp(ip, &elf, interp, sizeof(interp)) < 0)
    goto bad;
  if (load_image(pgdir, ip, &elf, base, &max_sz, &phdr_va) < 0)
    goto bad;
  iunlockput(ip);
  ip = 0;

  // The interpreter, at the bottom of the mmap area. It must itself be
  // position independent and must not ask for an interpreter of its own.
  if (interp[0]) {
    struct elfhdr ielf;
    uint64 ihdr;
    if ((ip = nameifollow(0, interp, &err)) == 0) {
      ret = -err;
      goto bad;
    }
    ilock(ip);
    if (readi(ip, 0, (uint64)&ielf, 0, sizeof(ielf)) != sizeof(ielf) ||
        ielf.magic != ELF_MAGIC || ielf.type != ET_DYN)
      goto bad;
    if (load_image(pgdir, ip, &ielf, INTERP_BASE, &interp_end, &ihdr) < 0)
      goto bad;
    interp_entry = INTERP_BASE + ielf.entry;
    iunlockput(ip);
    ip = 0;
  }
  end_op();
  in_op = 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 + 32;
  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;
    // Where the program is, for the dynamic linker -- and for a static musl
    // program too, which finds its TLS template through AT_PHDR.
    uint64 where[] = {
      3, phdr_va,                   // AT_PHDR
      4, sizeof(struct proghdr),    // AT_PHENT
      5, elf.phnum,                 // AT_PHNUM
      7, interp[0] ? INTERP_BASE : 0, // AT_BASE: the interpreter's load address
      9, base + elf.entry,          // AT_ENTRY: the program's, not the interpreter's
    };
    for (int k = 0; k < (int)NELEM(where); k++)
      if (tbl_put(&t, where[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));
  if (strace_wanted(myproc()->name))
    myproc()->strace = 1;
  safestrcpy(myproc()->exe, path, sizeof(myproc()->exe));

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

  // Auto-select syscall ABI: musl binaries have entry >= 0x100000;
  // SchOS-native binaries (init) are linked at 0x4000.
  if (base + 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;
  }

  // A new image is a new address space: a fresh mm, and the old one dropped
  // at the end. mm_put() detaches shared segments before freeing anything,
  // and does nothing at all while a vfork parent is still using it.
  if ((newmm = mm_alloc()) == 0)
    goto bad;
  proc_exec_threads();
  newmm->pgdir = pgdir;
  newmm->sz = sz;
  oldmm = myproc()->mm;
  myproc()->mm = newmm;
  myproc()->pgdir = pgdir;
  myproc()->sz = sz;
  pgdir = 0;                     // newmm owns it now; not bad:'s to free

  // 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.
  // Ignored signals stay ignored, and the mask is inherited (execve(2)):
  // neither points into the old image. `nohup` and every program that ignores
  // SIGPIPE before running a helper depend on it.
  for (int i = 0; i < 32; i++)
    if (myproc()->sighand->handler[i] != 1)
      myproc()->sighand->handler[i] = 0;
  memset(myproc()->sighand->restorer, 0, sizeof(myproc()->sighand->restorer));
  myproc()->fs_base          = 0;
  myproc()->pending_signals  = 0;
  myproc()->in_signal        = 0;
  myproc()->sigsusp_active   = 0;
  myproc()->mm->mmap_next        = 0;

  // The interpreter is ordinary memory to everything after this: fork copies
  // it, exit frees it. mmap carries on above it.
  if (interp[0]) {
    myproc()->mm->vma[0].addr    = INTERP_BASE;
    myproc()->mm->vma[0].length  = PGROUNDUP(interp_end);
    myproc()->mm->vma[0].used    = 1;
    myproc()->mm->vma[0].no_free = 0;
    myproc()->mm->vma[0].shmid   = -1;
    myproc()->mm->mmap_next      = INTERP_BASE + PGROUNDUP(interp_end);
  }

  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;
  uint64 entry = interp[0] ? interp_entry : base + elf.entry;
  myproc()->tf->rip = entry;
  myproc()->tf->rcx = entry;
  myproc()->tf->rsp = sp;

  // syscall_entry saves user RSP at kstack+KSTACKSIZE-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 + KSTACKSIZE - 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());
  // A new program starts with the FPU as the ABI promises it, not with the
  // old one's registers.
  memmove(myproc()->fpu, fpu_default, sizeof(myproc()->fpu));
  fxrstor(myproc()->fpu);
  mm_put(oldmm);
  proc_exec_done();              // a vfork parent may have its memory back
  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);
  if(in_op)
    end_op();
  return ret;
}