// Per-worker compile, evaluate and JS heap bench for a runtime.js bundle, in a real
// Chromium module worker: sketches/lazy-builtins.md §1's method. Spawns idle module
// workers on the target bundle (they never receive `init`, same as the kernel's spare
// worker), reads a trace for v8.compileModule/v8.evaluateModule per worker thread, and
// reads post-GC JS heap over the DevTools protocol. The first worker of each run is
// dropped (its compile/evaluate/heap is warm-up noise), matching the doc's method.
//
// Usage: node scripts/bench-worker.mjs [path/to/runtime.js] [path/to/variant.js] [runs]
//   path/to/runtime.js   defaults to dist/runtime/runtime.js
//   path/to/variant.js   optional: a second bundle to bench against, for comparison
//   runs                 rounds per variant, default 3 (>=9 worker samples per variant)
import { chromium } from "@playwright/test";
import { createServer } from "node:http";
import { readFileSync } from "node:fs";
import { headersFor } from "../src/node/server.ts";

const [, , argA, argB, argRuns] = process.argv;
const VARIANTS = { a: argA ?? "dist/runtime/runtime.js" };
if (argB) VARIANTS.b = argB;
const RUNS = Number(argRuns ?? 3);
const WORKERS = 9; // >= 9 samples per variant, per the doc
const GAP = 300;

let current = "a";

const PROBE_HTML = `<!doctype html><meta charset="utf-8"><title>probe</title><p id="exit">…</p>
<script type="module" src="/probe.js"></script>`;

const PROBE_JS = `
const q = new URLSearchParams(location.search);
const n = Number(q.get("n"));
const gap = Number(q.get("gap"));
window.workers = [];
for (let i = 0; i < n; i++) {
  workers.push(new Worker("/runtime.js", { type: "module", name: "probe-" + i }));
  await new Promise((r) => setTimeout(r, gap));
}
document.getElementById("exit").textContent = "0";
`;

// COOP/COEP/CORP must be on every response, not just probe.html's: COEP require-corp
// gates every subresource, workers scripts included, on a CORP header (DESIGN.md; the
// same requirement src/node/server.ts's own headersFor documents).
const server = createServer((req, res) => {
  const path = new URL(req.url ?? "/", "http://localhost").pathname;
  const isolated = headersFor(path, { isolate: () => true });
  if (path === "/probe.html") {
    res.writeHead(200, { "Content-Type": "text/html; charset=utf-8", ...isolated }).end(PROBE_HTML);
  } else if (path === "/probe.js") {
    res.writeHead(200, { "Content-Type": "text/javascript; charset=utf-8", ...isolated }).end(PROBE_JS);
  } else if (path === "/runtime.js") {
    res.writeHead(200, { "Content-Type": "text/javascript; charset=utf-8", ...isolated }).end(readFileSync(VARIANTS[current]));
  } else {
    res.writeHead(404).end();
  }
});
await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve));
const base = `http://127.0.0.1:${server.address().port}`;

// A raw CDP connection over --remote-debugging-port, alongside Playwright's own: Playwright's
// public CDPSession.send() has no way to address another target's session, and getting a
// worker's post-GC heap (Runtime.getHeapUsage) and per-thread compile/evaluate durations
// needs exactly that. flatten:true puts every event and reply on one socket, tagged by
// sessionId.
const cdpPort = 9340 + Math.floor(Math.random() * 500);
const browser = await chromium.launch({ args: [`--remote-debugging-port=${cdpPort}`] });
const { webSocketDebuggerUrl } = await (await fetch(`http://127.0.0.1:${cdpPort}/json/version`)).json();
const ws = new WebSocket(webSocketDebuggerUrl);
await new Promise((resolve) => (ws.onopen = resolve));

let id = 0;
const waiting = new Map();
ws.onmessage = (ev) => {
  const m = JSON.parse(ev.data);
  if (m.id && waiting.has(m.id)) {
    const w = waiting.get(m.id);
    waiting.delete(m.id);
    m.error ? w.reject(new Error(JSON.stringify(m.error))) : w.resolve(m.result);
  }
};
// A command to a target that never answers (closed mid-flight, or a domain that isn't
// enabled yet) must not hang the whole bench: 3s is generous for a local, same-process call.
function send(method, params = {}, sessionId) {
  const msg = { id: ++id, method, params, ...(sessionId ? { sessionId } : {}) };
  ws.send(JSON.stringify(msg));
  return new Promise((resolve, reject) => {
    const timer = setTimeout(() => {
      waiting.delete(msg.id);
      reject(new Error(`${method} timed out`));
    }, 3000);
    waiting.set(msg.id, {
      resolve: (v) => { clearTimeout(timer); resolve(v); },
      reject: (e) => { clearTimeout(timer); reject(e); },
    });
  });
}

async function runOnce(label) {
  current = label;
  const ctx = await browser.newContext();
  const page = await ctx.newPage();
  await browser.startTracing(page, {
    categories: ["v8", "devtools.timeline", "disabled-by-default-v8.compile", "blink", "loading", "__metadata"],
  });
  await page.goto(`${base}/probe.html?n=${WORKERS}&gap=${GAP}`);
  await page.waitForFunction(() => document.getElementById("exit")?.textContent === "0", null, { timeout: 60_000, polling: 5 });
  await new Promise((r) => setTimeout(r, 500)); // let the last worker's module finish evaluating

  // Attach to each worker target directly: runtime.js's own evaluation already finished
  // by the time probe.js's loop moves on, so there is no risk of attaching mid-evaluate
  // and skewing the trace (the trace, not this attach, is what timed it).
  const { targetInfos } = await send("Target.getTargets", {});
  const heaps = [];
  for (const t of targetInfos) {
    if (t.type !== "worker") continue;
    try {
      const { sessionId } = await send("Target.attachToTarget", { targetId: t.targetId, flatten: true });
      // Chromium won't answer any further command on a freshly attached worker session
      // until this is sent, even with waitForDebuggerOnStart left at its default (false).
      await send("Runtime.runIfWaitingForDebugger", {}, sessionId).catch(() => {});
      await send("Runtime.enable", {}, sessionId);
      await send("HeapProfiler.enable", {}, sessionId);
      await send("HeapProfiler.collectGarbage", {}, sessionId).catch(() => {});
      const h = await send("Runtime.getHeapUsage", {}, sessionId);
      heaps.push(Math.round(h.usedSize / 1024));
    } catch {
      // A worker that doesn't answer in time is dropped from this round's heap sample.
    }
  }

  const buf = await browser.stopTracing();
  const events = JSON.parse(buf.toString()).traceEvents;
  const compiles = events
    .filter((e) => e.ph === "X" && e.name === "v8.compileModule" && e.args?.data?.url?.endsWith("runtime.js"))
    .sort((a, b) => a.ts - b.ts);
  const tids = new Set(compiles.map((e) => e.tid));
  const evaluates = events
    .filter((e) => e.ph === "X" && e.name === "v8.evaluateModule" && tids.has(e.tid) && e.args?.worldId !== 0)
    .sort((a, b) => a.ts - b.ts);
  await ctx.close();
  // Drop the first worker of this run (warm-up noise), per the doc.
  return {
    compile: compiles.slice(1).map((e) => e.dur / 1000),
    evaluate: evaluates.slice(1).map((e) => e.dur / 1000),
    heap: heaps.slice(1),
  };
}

function median(a) {
  if (!a.length) return NaN;
  const s = [...a].sort((x, y) => x - y);
  const n = s.length;
  return n % 2 ? s[(n - 1) / 2] : (s[n / 2 - 1] + s[n / 2]) / 2;
}
function p10(a) {
  if (!a.length) return NaN;
  const s = [...a].sort((x, y) => x - y);
  return s[Math.max(0, Math.floor(s.length * 0.1))];
}

const results = {};
for (const label of Object.keys(VARIANTS)) results[label] = { compile: [], evaluate: [], heap: [] };
// Interleaved rounds, matching the doc's sweep: compare variants under the same conditions.
for (let r = 0; r < RUNS; r++) {
  for (const label of Object.keys(VARIANTS)) {
    const out = await runOnce(label);
    results[label].compile.push(...out.compile);
    results[label].evaluate.push(...out.evaluate);
    results[label].heap.push(...out.heap);
  }
}

for (const [label, path] of Object.entries(VARIANTS)) {
  const r = results[label];
  console.log(
    `${label} (${path})  compile med ${median(r.compile).toFixed(2)}ms (p10 ${p10(r.compile).toFixed(2)})` +
      `  evaluate med ${median(r.evaluate).toFixed(2)}ms (p10 ${p10(r.evaluate).toFixed(2)})` +
      `  heap med ${median(r.heap)}KB  n=${r.compile.length}`,
  );
}

ws.close();
await browser.close();
server.close();