#!/usr/bin/env node
// wasm-export-consult: export a main's imported xmod dispatch consult.
//
// BRINTOS-WORKAROUND: linux-in-the-browser#199 remove — the released frame binds the kernel-embedded runtime's env.__hwjs_xmod_dispatch_check from the main's exports, and a main only imports it.
//
// The frame late-binds the runtime's env.* imports three ways: a kernel-owned
// set against the vmlinux exports, a runtime-op set against module #0, and
// everything else against THE MAIN's exports. The consult is in the last group,
// but the pass only ever emits it as an env function IMPORT, so no main defines
// it and no --export flag can surface it (wasm-ld drops a hard --export of an
// import-only symbol without a diagnostic). The runtime's first call to it, in
// libc's constructors, then aborts the process before main.
//
// The export names a THUNK appended as the last defined function: it forwards
// its parameters to the consult's import and returns what that returns. An
// export record naming the import's funcidx directly is valid wasm and runs,
// but LLVM's object reader assumes every function export has a body, so
// llvm-nm crashes on such a module, and the build reads mains with it.
//
// Appending the function moves no existing index, and the code section's count
// keeps its encoded width, so every existing body stays at its offset within the
// section and the DWARF that addresses them stays correct. Only a count at a
// LEB128 width boundary (127, 16383, ...) has to widen; the output says so.
//
// Usage: wasm-export-consult.mjs <main.wasm>
//
// Prints one line and exits 0:
//   exported <thunk> -> <import>   the file was rewritten in place (inode and
//                                  mode kept); both are funcidx. A trailing
//                                  "(code count widened ...)" reports the
//                                  boundary case above.
//   already-exported               the main already exports a function by that name
//   no-import                      the main does not import the consult; untouched
// Exits 2 on anything it cannot read, parse or re-validate, before touching the
// file: the rewrite is compiled and re-read first. A failed write is exit 2 as
// well, and the caller discards the file.
import fs from 'node:fs';

const SYM = '__hwjs_xmod_dispatch_check';

function die(msg) {
  console.error(`wasm-export-consult: ${msg}`);
  process.exit(2);
}

function uleb(b, p) {
  let r = 0, s = 0, x;
  do {
    if (p >= b.length) die('truncated LEB128');
    x = b[p++];
    r += (x & 0x7f) * 2 ** s;
    s += 7;
  } while (x & 0x80);
  return [r, p];
}

function leb(n, width = 0) {
  const out = [];
  do {
    out.push(n % 128);
    n = Math.floor(n / 128);
  } while (n);
  while (out.length < width) out.push(0);
  for (let i = 0; i < out.length - 1; i++) out[i] |= 0x80;
  return Buffer.from(out);
}

// Each section with its raw bytes, header included, so an untouched one is
// copied exactly as the linker wrote it.
function sections(b) {
  const out = [];
  let p = 8;
  while (p < b.length) {
    const start = p, id = b[p++];
    let size;
    [size, p] = uleb(b, p);
    out.push({ id, raw: b.subarray(start, p + size), body: b.subarray(p, p + size) });
    p += size;
  }
  return out;
}

// Function imports in order, with their type index; the other kinds are skipped
// by their encodings, which is all this needs from them.
function functionImports(body) {
  let [n, p] = uleb(body, 0);
  const fns = [];
  const skipLimits = () => {
    const flags = body[p++];
    [, p] = uleb(body, p);
    if (flags & 1) [, p] = uleb(body, p);
  };
  for (let i = 0; i < n; i++) {
    let len;
    [len, p] = uleb(body, p);
    const mod = body.toString('utf8', p, p + len); p += len;
    [len, p] = uleb(body, p);
    const name = body.toString('utf8', p, p + len); p += len;
    const kind = body[p++];
    if (kind === 0) {
      let t;
      [t, p] = uleb(body, p);
      fns.push({ mod, name, type: t });
    } else if (kind === 1) { p++; skipLimits(); }
    else if (kind === 2) skipLimits();
    else if (kind === 3) p += 2;
    else die(`import '${mod}.${name}' has kind ${kind}, which this rewrite does not parse`);
  }
  return fns;
}

function paramCount(body, want) {
  let [n, p] = uleb(body, 0);
  if (want >= n) die(`type index ${want} is out of range`);
  for (let i = 0; i < n; i++) {
    if (body[p++] !== 0x60) die(`type ${i} is not a plain function type`);
    let np, nr;
    [np, p] = uleb(body, p);
    if (i === want) return np;
    p += np;
    [nr, p] = uleb(body, p);
    p += nr;
  }
  return -1;
}

const args = process.argv.slice(2);
if (args.length !== 1) die('usage: wasm-export-consult.mjs <main.wasm>');
const file = args[0];

let bytes, mod;
try {
  bytes = fs.readFileSync(file);
  mod = new WebAssembly.Module(bytes);
} catch (e) {
  die(`cannot read ${file} as a wasm module: ${e.message}`);
}

const secs = sections(bytes);
const sec = (id) => secs.find((s) => s.id === id);
const imports = sec(2) ? functionImports(sec(2).body) : [];
const consultIdx = imports.findIndex((f) => f.mod === 'env' && f.name === SYM);
if (consultIdx < 0) {
  console.log('no-import');
  process.exit(0);
}

const clash = WebAssembly.Module.exports(mod).find((e) => e.name === SYM);
if (clash) {
  if (clash.kind !== 'function') die(`${file} exports '${SYM}' as a ${clash.kind}, not a function`);
  console.log('already-exported');
  process.exit(0);
}

if (!sec(7)) die(`${file} imports '${SYM}' but has no export section to add it to`);
if (!sec(3) || !sec(10)) die(`${file} imports '${SYM}' but defines no function to append the thunk after`);

const type = imports[consultIdx].type;
const nparams = paramCount(sec(1).body, type);

const [nfuncs, fnFirst] = uleb(sec(3).body, 0);
const thunkIdx = imports.length + nfuncs;

const [ncode, codeFirst] = uleb(sec(10).body, 0);
if (ncode !== nfuncs) die(`function count ${nfuncs} and code count ${ncode} disagree`);
const newCount = leb(ncode + 1, codeFirst);
const widened = newCount.length - codeFirst;

const insns = [Buffer.from([0x00])];
for (let i = 0; i < nparams; i++) insns.push(Buffer.from([0x20]), leb(i));
insns.push(Buffer.from([0x10]), leb(consultIdx), Buffer.from([0x0b]));
const thunk = Buffer.concat(insns);

const [nexp, expFirst] = uleb(sec(7).body, 0);
const name = Buffer.from(SYM);

const payload = {
  3: Buffer.concat([leb(nfuncs + 1), sec(3).body.subarray(fnFirst), leb(type)]),
  10: Buffer.concat([newCount, sec(10).body.subarray(codeFirst), leb(thunk.length), thunk]),
  7: Buffer.concat([leb(nexp + 1), sec(7).body.subarray(expFirst),
    leb(name.length), name, Buffer.from([0x00]), leb(thunkIdx)]),
};
const out = Buffer.concat([
  bytes.subarray(0, 8),
  ...secs.map((s) => (s.id in payload
    ? Buffer.concat([Buffer.from([s.id]), leb(payload[s.id].length), payload[s.id]])
    : s.raw)),
]);

let check;
try {
  check = new WebAssembly.Module(out);
} catch (e) {
  die(`the rewritten module does not validate: ${e.message}`);
}
const got = WebAssembly.Module.exports(check).filter((e) => e.name === SYM && e.kind === 'function');
if (got.length !== 1) die(`the rewritten module carries ${got.length} function export(s) named '${SYM}', expected 1`);

try {
  fs.writeFileSync(file, out);
} catch (e) {
  die(`cannot write ${file}: ${e.message}`);
}
console.log(`exported ${thunkIdx} -> ${consultIdx}` +
  (widened ? ` (code count widened by ${widened} byte; code offsets after it moved)` : ''));