// OxygenFlow's WASM tier in a browser tab (oxygenflow#80): this worker
// unpacks a repository's zip archive into memory and runs a job's WASI
// preview 1 modules over it, as `flow` runs them with wasmtime: the files at
// /src (and as the current directory), the job's arguments and environment,
// stdout and stderr streamed back line by line, and the exit code.
//
// Messages in: { archive: ArrayBuffer, exec: [module, ...args], env: {k: v} }.
// Messages out: { kind: "note"|"out"|"err", text }, then
// { kind: "done", code, ms } or { kind: "failed", text }.
"use strict";

const E = { SUCCESS: 0, ACCES: 2, BADF: 8, EXIST: 20, INVAL: 28, ISDIR: 31, NOENT: 44, NOSYS: 52, NOTDIR: 54, NOTEMPTY: 55, NOTSUP: 58 };
const FT = { UNKNOWN: 0, CHAR: 2, DIR: 3, FILE: 4 };
const O = { CREAT: 1, DIRECTORY: 2, EXCL: 4, TRUNC: 8 };
const FD_APPEND = 1;
const RIGHTS = 0x1fffffffn;

class Exit {
  constructor(code) { this.code = code; }
}

// --- The archive -----------------------------------------------------------

// A zip's files by path, the first directory (the archive's prefix) taken
// off. Read through the central directory, so sizes written after the data
// are known; stored and deflated entries; symbolic links are left out.
async function unzip(buf) {
  const b = new DataView(buf);
  const bytes = new Uint8Array(buf);
  let end = -1;
  for (let i = buf.byteLength - 22; i >= Math.max(0, buf.byteLength - 65557); i--) {
    if (b.getUint32(i, true) === 0x06054b50) { end = i; break; }
  }
  if (end < 0) throw new Error("the archive is not a zip file");
  const count = b.getUint16(end + 10, true);
  let at = b.getUint32(end + 16, true);
  if (at === 0xffffffff) throw new Error("the archive is too large for this page (zip64)");
  const files = new Map();
  const dec = new TextDecoder();
  for (let n = 0; n < count; n++) {
    if (b.getUint32(at, true) !== 0x02014b50) throw new Error("the archive's directory is damaged");
    const method = b.getUint16(at + 10, true);
    const csize = b.getUint32(at + 20, true);
    const nlen = b.getUint16(at + 28, true), xlen = b.getUint16(at + 30, true), clen = b.getUint16(at + 32, true);
    const mode = b.getUint32(at + 38, true) >>> 16;
    const local = b.getUint32(at + 42, true);
    const name = dec.decode(bytes.subarray(at + 46, at + 46 + nlen));
    at += 46 + nlen + xlen + clen;
    if (name.endsWith("/") || (mode & 0o170000) === 0o120000) continue;
    const data = local + 30 + b.getUint16(local + 26, true) + b.getUint16(local + 28, true);
    const raw = bytes.subarray(data, data + csize);
    let content;
    if (method === 0) content = raw.slice();
    else if (method === 8) content = new Uint8Array(await new Response(new Blob([raw]).stream().pipeThrough(new DecompressionStream("deflate-raw"))).arrayBuffer());
    else throw new Error(`${name}: compression method ${method} is not read here`);
    const slash = name.indexOf("/");
    files.set(slash < 0 ? name : name.slice(slash + 1), content);
  }
  return files;
}

// --- The file system ---------------------------------------------------------

class Fs {
  constructor(files) {
    this.files = files; // path -> Uint8Array
    this.dirs = new Set([""]);
    for (const p of files.keys()) this.addParents(p);
  }
  addParents(p) {
    const parts = p.split("/");
    for (let i = 1; i < parts.length; i++) this.dirs.add(parts.slice(0, i).join("/"));
  }
  // `rel` against directory `dir`, or null when it leaves the root.
  join(dir, rel) {
    const out = rel.startsWith("/") ? [] : dir ? dir.split("/") : [];
    for (const part of rel.split("/")) {
      if (part === "" || part === ".") continue;
      if (part === "..") { if (!out.length) return null; out.pop(); } else out.push(part);
    }
    return out.join("/");
  }
  kind(p) {
    if (this.files.has(p)) return FT.FILE;
    if (this.dirs.has(p)) return FT.DIR;
    return null;
  }
  list(dir) {
    const prefix = dir ? dir + "/" : "";
    const names = new Set();
    for (const p of [...this.files.keys(), ...this.dirs]) {
      if (p !== dir && p.startsWith(prefix)) names.add(p.slice(prefix.length).split("/")[0]);
    }
    return [...names].sort();
  }
}

// The modules a job's first `exec` word names: one path, or a `*` pattern
// in its file name, as `flow` resolves them.
function modules(fs, program) {
  const rel = program.replace(/^\/src\//, "").replace(/^\.\//, "");
  if (!rel || rel.startsWith("/") || rel.split("/").includes("..")) throw new Error(`module ${program}: give a path inside the job's files`);
  const cut = rel.lastIndexOf("/");
  const dir = cut < 0 ? "" : rel.slice(0, cut), last = rel.slice(cut + 1);
  if (dir.includes("*")) throw new Error(`module ${program}: only the file name may hold *`);
  if (!last.includes("*")) {
    if (!fs.files.has(rel)) throw new Error(`module ${rel} is not among the repository's files (a build makes it?)`);
    return [rel];
  }
  const re = new RegExp("^" + last.split("*").map((s) => s.replace(/[.+?^${}()|[\]\\]/g, "\\$&")).join(".*") + "$");
  const found = fs.list(dir).filter((n) => re.test(n) && fs.files.has(dir ? `${dir}/${n}` : n)).map((n) => (dir ? `${dir}/${n}` : n));
  if (!found.length) throw new Error(`no module matches ${program}`);
  return found;
}

// --- WASI preview 1 ------------------------------------------------------------

function wasi(fs, argv, env, say) {
  let memory;
  const view = () => new DataView(memory.buffer);
  const u8 = () => new Uint8Array(memory.buffer);
  const enc = new TextEncoder(), dec = new TextDecoder();
  const str = (ptr, len) => dec.decode(u8().subarray(ptr, ptr + len));
  // Open files and directories: 0-2 the streams, 3 /src, 4 ".".
  const fds = new Map([
    [0, { stream: "in" }],
    [1, { stream: "out", partial: "" }],
    [2, { stream: "err", partial: "" }],
    [3, { dir: "", preopen: "/src" }],
    [4, { dir: "", preopen: "." }],
  ]);
  let next = 5;
  const emit = (fd, text) => {
    const f = fds.get(fd);
    f.partial += text;
    let nl;
    while ((nl = f.partial.indexOf("\n")) >= 0) {
      say(f.stream, f.partial.slice(0, nl));
      f.partial = f.partial.slice(nl + 1);
    }
  };
  const flush = () => { for (const fd of [1, 2]) { const f = fds.get(fd); if (f.partial) { say(f.stream, f.partial); f.partial = ""; } } };
  const filestat = (ptr, kind, size) => {
    const v = view();
    for (let i = 0; i < 64; i += 8) v.setBigUint64(ptr + i, 0n, true);
    v.setUint8(ptr + 16, kind);
    v.setBigUint64(ptr + 24, 1n, true);
    v.setBigUint64(ptr + 32, BigInt(size), true);
  };
  const iovs = (ptr, n) => {
    const v = view(), out = [];
    for (let i = 0; i < n; i++) out.push([v.getUint32(ptr + i * 8, true), v.getUint32(ptr + i * 8 + 4, true)]);
    return out;
  };
  const strings = (list) => list.map((s) => enc.encode(s + "\0"));
  const sizes = (list, countPtr, sizePtr) => {
    const v = view(), all = strings(list);
    v.setUint32(countPtr, all.length, true);
    v.setUint32(sizePtr, all.reduce((n, s) => n + s.length, 0), true);
    return E.SUCCESS;
  };
  const fill = (list, ptrs, buf) => {
    const v = view(), m = u8();
    for (const s of strings(list)) { v.setUint32(ptrs, buf, true); ptrs += 4; m.set(s, buf); buf += s.length; }
    return E.SUCCESS;
  };
  const envList = Object.entries(env).map(([k, v]) => `${k}=${v}`);
  const imports = {
    args_sizes_get: (c, s) => sizes(argv, c, s),
    args_get: (p, b) => fill(argv, p, b),
    environ_sizes_get: (c, s) => sizes(envList, c, s),
    environ_get: (p, b) => fill(envList, p, b),
    clock_time_get: (_id, _precision, out) => { view().setBigUint64(out, BigInt(Math.round(performance.timeOrigin + performance.now())) * 1000000n, true); return E.SUCCESS; },
    clock_res_get: (_id, out) => { view().setBigUint64(out, 1000n, true); return E.SUCCESS; },
    random_get: (ptr, len) => { crypto.getRandomValues(u8().subarray(ptr, ptr + len)); return E.SUCCESS; },
    sched_yield: () => E.SUCCESS,
    proc_exit: (code) => { throw new Exit(code); },
    fd_prestat_get: (fd, ptr) => {
      const f = fds.get(fd);
      if (!f || !f.preopen) return E.BADF;
      const v = view();
      v.setUint8(ptr, 0);
      v.setUint32(ptr + 4, enc.encode(f.preopen).length, true);
      return E.SUCCESS;
    },
    fd_prestat_dir_name: (fd, ptr, len) => {
      const f = fds.get(fd);
      if (!f || !f.preopen) return E.BADF;
      u8().set(enc.encode(f.preopen).subarray(0, len), ptr);
      return E.SUCCESS;
    },
    fd_fdstat_get: (fd, ptr) => {
      const f = fds.get(fd);
      if (!f) return E.BADF;
      const v = view();
      v.setUint8(ptr, f.stream ? FT.CHAR : f.dir !== undefined ? FT.DIR : FT.FILE);
      v.setUint16(ptr + 2, f.append ? FD_APPEND : 0, true);
      v.setBigUint64(ptr + 8, RIGHTS, true);
      v.setBigUint64(ptr + 16, RIGHTS, true);
      return E.SUCCESS;
    },
    fd_fdstat_set_flags: () => E.SUCCESS,
    fd_filestat_get: (fd, ptr) => {
      const f = fds.get(fd);
      if (!f) return E.BADF;
      if (f.stream) filestat(ptr, FT.CHAR, 0);
      else if (f.dir !== undefined) filestat(ptr, FT.DIR, 0);
      else filestat(ptr, FT.FILE, fs.files.get(f.path).length);
      return E.SUCCESS;
    },
    path_filestat_get: (fd, _flags, p, len, ptr) => {
      const d = fds.get(fd);
      if (!d || d.dir === undefined) return E.BADF;
      const path = fs.join(d.dir, str(p, len));
      if (path === null) return E.ACCES;
      const kind = fs.kind(path);
      if (kind === null) return E.NOENT;
      filestat(ptr, kind, kind === FT.FILE ? fs.files.get(path).length : 0);
      return E.SUCCESS;
    },
    path_open: (fd, _dirflags, p, len, oflags, _rb, _ri, fdflags, out) => {
      const d = fds.get(fd);
      if (!d || d.dir === undefined) return E.BADF;
      const path = fs.join(d.dir, str(p, len));
      if (path === null) return E.ACCES;
      let kind = fs.kind(path);
      if (oflags & O.DIRECTORY) {
        if (kind === null) return E.NOENT;
        if (kind !== FT.DIR) return E.NOTDIR;
      }
      if (kind === null) {
        if (!(oflags & O.CREAT)) return E.NOENT;
        const parent = path.includes("/") ? path.slice(0, path.lastIndexOf("/")) : "";
        if (fs.kind(parent) !== FT.DIR) return E.NOENT;
        fs.files.set(path, new Uint8Array(0));
        kind = FT.FILE;
      } else if (oflags & O.EXCL) return E.EXIST;
      if (kind === FT.DIR) fds.set(next, { dir: path, entries: null });
      else {
        if (oflags & O.TRUNC) fs.files.set(path, new Uint8Array(0));
        fds.set(next, { path, pos: 0, append: !!(fdflags & FD_APPEND) });
      }
      view().setUint32(out, next, true);
      next += 1;
      return E.SUCCESS;
    },
    path_create_directory: (fd, p, len) => {
      const d = fds.get(fd);
      if (!d || d.dir === undefined) return E.BADF;
      const path = fs.join(d.dir, str(p, len));
      if (path === null) return E.ACCES;
      if (fs.kind(path) !== null) return E.EXIST;
      fs.dirs.add(path);
      return E.SUCCESS;
    },
    path_unlink_file: (fd, p, len) => {
      const d = fds.get(fd);
      if (!d || d.dir === undefined) return E.BADF;
      const path = fs.join(d.dir, str(p, len));
      if (path === null || !fs.files.has(path)) return E.NOENT;
      fs.files.delete(path);
      return E.SUCCESS;
    },
    fd_close: (fd) => {
      if (!fds.has(fd) || fd < 3) return fds.has(fd) ? E.SUCCESS : E.BADF;
      fds.delete(fd);
      return E.SUCCESS;
    },
    fd_read: (fd, ptr, n, out) => {
      const f = fds.get(fd);
      if (!f) return E.BADF;
      if (f.dir !== undefined) return E.ISDIR;
      let total = 0;
      if (!f.stream) {
        const data = fs.files.get(f.path), m = u8();
        for (const [buf, len] of iovs(ptr, n)) {
          const k = Math.min(len, data.length - f.pos);
          if (k <= 0) break;
          m.set(data.subarray(f.pos, f.pos + k), buf);
          f.pos += k;
          total += k;
        }
      }
      view().setUint32(out, total, true);
      return E.SUCCESS;
    },
    fd_write: (fd, ptr, n, out) => {
      const f = fds.get(fd);
      if (!f) return E.BADF;
      if (f.dir !== undefined) return E.ISDIR;
      let total = 0;
      const m = u8();
      for (const [buf, len] of iovs(ptr, n)) {
        const chunk = m.slice(buf, buf + len);
        if (f.stream === "out" || f.stream === "err") emit(fd, dec.decode(chunk, { stream: true }));
        else if (!f.stream) {
          const old = fs.files.get(f.path);
          const at = f.append ? old.length : f.pos;
          const grown = new Uint8Array(Math.max(old.length, at + len));
          grown.set(old);
          grown.set(chunk, at);
          fs.files.set(f.path, grown);
          f.pos = at + len;
        }
        total += len;
      }
      view().setUint32(out, total, true);
      return E.SUCCESS;
    },
    fd_seek: (fd, offset, whence, out) => {
      const f = fds.get(fd);
      if (!f || f.stream || f.dir !== undefined) return E.BADF;
      const size = fs.files.get(f.path).length;
      const base = whence === 0 ? 0 : whence === 1 ? f.pos : size;
      const pos = base + Number(offset);
      if (pos < 0) return E.INVAL;
      f.pos = pos;
      view().setBigUint64(out, BigInt(pos), true);
      return E.SUCCESS;
    },
    fd_tell: (fd, out) => {
      const f = fds.get(fd);
      if (!f || f.stream || f.dir !== undefined) return E.BADF;
      view().setBigUint64(out, BigInt(f.pos), true);
      return E.SUCCESS;
    },
    fd_readdir: (fd, buf, len, cookie, out) => {
      const d = fds.get(fd);
      if (!d || d.dir === undefined) return E.BADF;
      if (!d.entries) d.entries = [".", "..", ...fs.list(d.dir)];
      const v = view(), m = u8();
      let used = 0;
      for (let i = Number(cookie); i < d.entries.length; i++) {
        const name = d.entries[i];
        const bytes = enc.encode(name);
        const child = name === "." || name === ".." ? d.dir : fs.join(d.dir, name);
        const head = new DataView(new ArrayBuffer(24));
        head.setBigUint64(0, BigInt(i + 1), true);
        head.setBigUint64(8, BigInt(i + 1), true);
        head.setUint32(16, bytes.length, true);
        head.setUint8(20, name === "." || name === ".." ? FT.DIR : fs.kind(child) ?? FT.UNKNOWN);
        const entry = new Uint8Array(24 + bytes.length);
        entry.set(new Uint8Array(head.buffer));
        entry.set(bytes, 24);
        const room = len - used;
        m.set(entry.subarray(0, Math.min(room, entry.length)), buf + used);
        used += Math.min(room, entry.length);
        if (used >= len) break;
      }
      v.setUint32(out, used, true);
      return E.SUCCESS;
    },
  };
  return { imports, flush, bind: (mem) => { memory = mem; } };
}

// --- A job -------------------------------------------------------------------------

async function run({ archive, exec, env }) {
  const say = (kind, text) => postMessage({ kind, text });
  const started = performance.now();
  const fs = new Fs(await unzip(archive));
  say("note", `${fs.files.size} files unpacked in ${Math.round(performance.now() - started)} ms`);
  const list = modules(fs, exec[0]);
  let first = 0;
  for (const path of list) {
    say("note", `running ${path}`);
    const name = path.slice(path.lastIndexOf("/") + 1);
    const w = wasi(fs, [name, ...exec.slice(1)], { ...env, PWD: "/src" }, say);
    const module = await WebAssembly.compile(fs.files.get(path));
    const imports = { wasi_snapshot_preview1: {} };
    // What the shim lacks answers ENOSYS rather than refusing the module.
    for (const imp of WebAssembly.Module.imports(module)) {
      if (imp.kind !== "function") continue;
      imports[imp.module] ??= {};
      if (imp.module === "wasi_snapshot_preview1" && w.imports[imp.name]) imports[imp.module][imp.name] = w.imports[imp.name];
      else imports[imp.module][imp.name] ??= () => E.NOSYS;
    }
    const instance = await WebAssembly.instantiate(module, imports);
    w.bind(instance.exports.memory);
    let code = 0;
    try {
      instance.exports._start();
    } catch (e) {
      if (e instanceof Exit) code = e.code;
      else {
        w.flush();
        say("err", `${path} trapped: ${e.message}`);
        code = 134;
      }
    }
    w.flush();
    if (code !== 0 && first === 0) first = code;
  }
  postMessage({ kind: "done", code: first, ms: Math.round(performance.now() - started) });
}

onmessage = (m) => {
  run(m.data).catch((e) => postMessage({ kind: "failed", text: e.message }));
};