// Throughput benchmark: serves the app with Vite, opens it in a headless
// Chromium with WebGPU, and runs the crash preset at several grid spacings,
// alone and with several concurrent runs, reporting steps per second and
// the wall time per second of audio. The numbers describe the machine's
// GPU as seen through Chromium; the kernel is latency-bound on small grids,
// so concurrency and finer grids are the interesting axes.
//
//   node scripts/bench.ts [--quick] [--preset crash-16] [--gpu-flags "..."] [--browsers N]
//
// --browsers N spreads the concurrent runs over N independent Chromium
// instances (each with its own GPU process), which tells whether a plateau
// in total throughput is the GPU's or the browser's.

import { fileURLToPath } from "node:url";
import { dirname, join } from "node:path";
import { createServer } from "vite";
import { chromium, type Browser } from "playwright";

const root = join(dirname(fileURLToPath(import.meta.url)), "..");
const args = process.argv.slice(2);
const quick = args.includes("--quick");
const preset = args.includes("--preset") ? args[args.indexOf("--preset") + 1] : "crash-16";
const extraFlags = args.includes("--gpu-flags") ? args[args.indexOf("--gpu-flags") + 1].split(" ") : [];
const nBrowsers = args.includes("--browsers") ? parseInt(args[args.indexOf("--browsers") + 1]) : 1;

const ARGS = ["--enable-unsafe-webgpu", "--enable-features=Vulkan", "--use-angle=vulkan", "--ignore-gpu-blocklist", "--no-sandbox", ...extraFlags];

interface Status { phase: string; totalSteps: number; dt: number; nx: number; ny: number; wmax: number; peakPressure: number }

interface Case { spacing: number; duration: number; concurrency: number }

const concOnly = args.includes("--conc");
const cases: Case[] = concOnly
  ? [{ spacing: 0.008, duration: 2, concurrency: 4 }, { spacing: 0.008, duration: 2, concurrency: 8 }]
  : quick
  ? [{ spacing: 0.008, duration: 1, concurrency: 1 }, { spacing: 0.004, duration: 0.25, concurrency: 1 }, { spacing: 0.008, duration: 1, concurrency: 4 }]
  : [
      { spacing: 0.008, duration: 2, concurrency: 1 },
      { spacing: 0.006, duration: 1, concurrency: 1 },
      { spacing: 0.004, duration: 0.5, concurrency: 1 },
      { spacing: 0.003, duration: 0.25, concurrency: 1 },
      { spacing: 0.008, duration: 2, concurrency: 2 },
      { spacing: 0.008, duration: 2, concurrency: 4 },
      { spacing: 0.008, duration: 2, concurrency: 8 },
      { spacing: 0.004, duration: 0.5, concurrency: 4 },
    ];

/** Runs one simulation to completion in a fresh page; resolves with its status and wall time. */
async function runOne(browser: Browser, url: string, c: Case): Promise<{ st: Status; wall: number; adapter: string }> {
  const page = await browser.newPage();
  const errors: string[] = [];
  page.on("pageerror", (e) => errors.push(e.message));
  page.on("console", (m) => { if (m.type() === "error") errors.push(m.text()); });
  await page.goto(`${url}?autostart=1&preset=${preset}&spacing=${c.spacing}&duration=${c.duration}`);
  const t0 = performance.now();
  let st: Status | null = null;
  for (let i = 0; i < 4000; i++) {
    await page.waitForTimeout(100);
    st = await page.evaluate<Status | null>("window.steelsim ? window.steelsim.status() : null");
    const noGpu = await page.evaluate<string | null>("window.steelsim ? window.steelsim.noGpu() : null");
    if (noGpu) throw new Error(noGpu);
    if (st && st.phase === "done") break;
  }
  const wall = (performance.now() - t0) / 1000;
  const adapter = await page.evaluate<string>(`navigator.gpu.requestAdapter().then(a => a ? (a.info ? a.info.vendor + " " + a.info.architecture + " " + a.info.device + " " + a.info.description : "adapter") : "none")`);
  await page.close();
  if (errors.length) throw new Error(errors.join("\n"));
  if (!st || st.phase !== "done") throw new Error("did not finish");
  return { st, wall, adapter };
}

async function main(): Promise<void> {
  const server = await createServer({ root, logLevel: "silent", server: { port: 0, strictPort: false } });
  await server.listen();
  const url = server.resolvedUrls!.local[0];
  const browsers = await Promise.all(Array.from({ length: nBrowsers }, () => chromium.launch({ headless: true, args: ARGS })));
  try {
    let adapter = "";
    console.log(`preset ${preset}, ${nBrowsers} browser instance${nBrowsers > 1 ? "s" : ""}`);
    console.log("spacing  grid       dt      conc  audio s  wall s   ksteps/s (each)  ksteps/s (total)  wall per audio s (each)");
    for (const c of cases) {
      const results = await Promise.all(Array.from({ length: c.concurrency }, (_, k) => runOne(browsers[k % nBrowsers], url, c)));
      adapter = results[0].adapter;
      const wall = Math.max(...results.map((r) => r.wall));
      const steps = results[0].st.totalSteps;
      const each = steps / wall / 1e3;
      const total = (steps * c.concurrency) / wall / 1e3;
      const st = results[0].st;
      console.log(
        `${(c.spacing * 1e3).toFixed(0).padStart(4)} mm  ${`${st.nx}x${st.ny}`.padEnd(9)} ${(st.dt * 1e6).toFixed(2).padStart(5)} us  ${String(c.concurrency).padStart(3)}   ${c.duration.toFixed(2).padStart(5)}   ${wall.toFixed(1).padStart(6)}   ${each.toFixed(0).padStart(8)}          ${total.toFixed(0).padStart(8)}          ${(wall / c.duration).toFixed(1).padStart(6)}x`,
      );
      if (!isFinite(st.wmax) || st.wmax > 0.05) console.log("  (diverged)");
    }
    console.log(`adapter: ${adapter}`);
  } finally {
    for (const b of browsers) await b.close();
    await server.close();
  }
}

void main();