/**
 * Render a scene to a .wav outside the browser, on whatever GPU this
 * machine has.
 *
 * The browser's Save scene button writes a scene .json; this script loads
 * it (or starts from a named preset), builds the same StudioSession the
 * page builds — same generator, same compiler, same solver — runs it flat
 * out with no display, and resamples the microphone's trace to a 48 kHz
 * .wav. Because the pipeline is identical, a scene renders the same
 * everywhere; what a bigger GPU buys is speed, and with --nx, a finer air
 * grid than the browser would enjoy: the audio bandwidth of the simulation
 * is c_air/(8h), so doubling nx doubles the honest bandwidth at 16x the
 * work per second of audio.
 *
 *   npm run render -- scene.json --seconds 10 --out clip.wav
 *   npm run render -- --preset drum --seconds 5 --nx 192
 *   npm run render -- --list-presets
 *
 * Options:
 *   --seconds N     duration of audio to render (default 5)
 *   --nx N          override the air grid (even; memory grows as nx^3)
 *   --mic x,y,z     override the microphone position, metres
 *   --out FILE      output path (default: the scene's or preset's name.wav)
 *
 * Runs through vite-node (which resolves numbl's compiler sources) on
 * desktop WebGPU via the optional `webgpu` package (Google Dawn). On a
 * headless Linux box Dawn needs working Vulkan; nodeWebGpu.ts explains
 * itself when it cannot load.
 */
import { readFile, writeFile } from 'node:fs/promises';
import { basename } from 'node:path';
import { installWebGpu, errMsg, NO_ADAPTER_HINT } from './nodeWebGpu.ts';
import { requestAcousticDevice } from '../src/device.ts';
import { StudioSession } from '../src/studio/session.ts';
import { parseSceneFile, makeSceneFile, type SceneFile } from '../src/studio/scenefile.ts';
import { PRESETS, cloneSpec, presetByKey } from '../src/studio/presets.ts';
import { specParams } from '../src/studio/session.ts';
import { stableDt } from '../src/grids.ts';
import { traceToWav, WAV_RATE } from '../src/audio/wav.ts';
import { C_AIR, CFL, DOMAIN_X, POOR_RESOLUTION, fmtTime } from '../src/units.ts';

function usage(): never {
  console.error(
    'usage: npm run render -- <scene.json | --preset KEY> ' +
      '[--seconds N] [--nx N] [--mic x,y,z] [--out FILE]\n' +
      '       npm run render -- --list-presets',
  );
  process.exit(2);
}

/* ---------------------------------------------------------------- args --- */

const args = process.argv.slice(2);
let scenePath: string | null = null;
let presetKey: string | null = null;
let seconds = 5;
let nxOverride: number | null = null;
let micOverride: { x: number; y: number; z: number } | null = null;
let outPath: string | null = null;

for (let i = 0; i < args.length; i++) {
  const a = args[i];
  const next = (): string => {
    const v = args[++i];
    if (v === undefined) usage();
    return v;
  };
  if (a === '--list-presets') {
    for (const p of PRESETS) console.log(`${p.key.padEnd(10)} ${p.blurb}`);
    process.exit(0);
  } else if (a === '--preset') presetKey = next();
  else if (a === '--seconds') seconds = Number(next());
  else if (a === '--nx') nxOverride = Number(next());
  else if (a === '--out') outPath = next();
  else if (a === '--mic') {
    const parts = next().split(',').map(Number);
    if (parts.length !== 3 || parts.some((v) => !Number.isFinite(v))) usage();
    micOverride = { x: parts[0], y: parts[1], z: parts[2] };
  } else if (a.startsWith('--')) usage();
  else if (scenePath === null) scenePath = a;
  else usage();
}
if ((scenePath === null) === (presetKey === null)) usage(); // exactly one
if (!(seconds > 0) || seconds > 600) {
  console.error('--seconds must be between 0 and 600');
  process.exit(2);
}

/* --------------------------------------------------------------- scene --- */

let file: SceneFile;
let name: string;
if (scenePath !== null) {
  const text = await readFile(scenePath, 'utf8').catch((e: unknown) => {
    console.error(`cannot read ${scenePath}: ${errMsg(e)}`);
    process.exit(1);
  });
  try {
    file = parseSceneFile(text);
  } catch (e) {
    console.error(errMsg(e));
    process.exit(1);
  }
  name = basename(scenePath).replace(/\.json$/i, '');
} else {
  const preset = presetByKey(presetKey!);
  if (!preset) {
    console.error(
      `no preset '${presetKey}' (one of: ${PRESETS.map((p) => p.key).join(', ')})`,
    );
    process.exit(1);
  }
  const spec = cloneSpec(preset.spec);
  file = makeSceneFile({ spec, ...specParams(spec), mic: preset.mic, nx: 96 });
  name = preset.key;
}

const nx = nxOverride ?? file.nx;
if (!Number.isInteger(nx) || nx < 16 || nx % 2 !== 0) {
  console.error(`--nx must be an even integer of at least 16 (got ${nx})`);
  process.exit(1);
}
const mic = micOverride ?? file.mic;
const out = outPath ?? `${name}.wav`;

/* -------------------------------------------------------------- render --- */

let device: GPUDevice;
try {
  const runtime = await installWebGpu();
  device = await requestAcousticDevice();
  console.log(`soundbox render — ${runtime}`);
} catch (e) {
  console.error(`render: ${errMsg(e)}\n${NO_ADAPTER_HINT}`);
  process.exit(1);
}

// Size the trace for the request before the session exists. The estimate
// assumes the medium tops out at the speed of air, which holds for every
// generated scene; a hand-edited faster medium shrinks dt, and the shortfall
// is reported below rather than silently truncated.
const dtEstimate = stableDt(DOMAIN_X / nx, C_AIR, CFL);
const capacity = Math.ceil((1.05 * seconds) / dtEstimate);
if (capacity > 300_000_000) {
  console.error(
    `${fmtTime(seconds)} at nx = ${nx} is ${capacity.toLocaleString()} samples of ` +
      `trace — more GPU memory than is sensible. Lower --seconds or --nx.`,
  );
  process.exit(1);
}

const t0 = performance.now();
const session = await StudioSession.create({
  device,
  spec: file.spec,
  params: file.params,
  sceneParams: file.sceneParams,
  modelSource: file.sources?.model,
  sceneSource: file.sources?.scene,
  nx,
  traceCapacity: capacity,
}).catch((e: unknown) => {
  console.error(`compile failed: ${errMsg(e)}`);
  process.exit(1);
});
const tCompiled = performance.now();

const { air } = session;
const fmax = C_AIR / (POOR_RESOLUTION * air.h);
const objects = file.spec.objects
  .map((o) => {
    const g = session.grids[o.id];
    return `${o.id} (${g.kind === 'line' ? `${g.n1} nodes` : `${g.n1}×${g.n2} nodes`})`;
  })
  .join(', ');
console.log(
  `air ${air.nx}×${air.ny}×${air.nz} (${(air.npts / 1e6).toFixed(2)} M cells, ` +
    `honest to ~${(fmax / 1000).toFixed(1)} kHz) · ${objects}\n` +
    `dt = ${fmtTime(session.dt)} (${Math.round(1 / session.dt / 1000)} kHz) · ` +
    `compiled in ${((tCompiled - t0) / 1000).toFixed(1)} s`,
);

session.setMic(mic.x, mic.y, mic.z);
session.excite();

const total = Math.ceil(seconds / session.dt);
if (total > session.recorder.capacity) {
  console.error(
    `the scene's medium lowered dt below the estimate; only ` +
      `${fmtTime(session.recorder.capacity * session.dt)} of trace fits. Rendering that.`,
  );
}
const steps = Math.min(total, session.recorder.capacity);
const batch = 2048;
let done = 0;
let lastReport = -1;
while (done < steps) {
  const n = Math.min(batch, steps - done);
  session.step(n);
  done += n;
  await session.sync();
  const pct = Math.floor((10 * done) / steps);
  if (pct > lastReport) {
    lastReport = pct;
    const dt = (performance.now() - tCompiled) / 1000;
    const rate = done / Math.max(dt, 1e-9);
    process.stderr.write(
      `  ${String(10 * pct).padStart(3)}%  ${(rate / 1000).toFixed(0)}k steps/s  ` +
        `${fmtTime(done * session.dt)} of audio\n`,
    );
  }
}
const tRun = performance.now();

const trace = await session.recorder.read();
let peak = 0;
for (const v of trace) peak = Math.max(peak, Math.abs(v));
if (!(peak > 0)) {
  console.error('warning: the microphone recorded silence — is every gain zero, or the mic inside a wall?');
}
const wav = traceToWav(trace, session.dt);
await writeFile(out, Buffer.from(await wav.arrayBuffer()));

const wall = (tRun - tCompiled) / 1000;
console.log(
  `${fmtTime(trace.length * session.dt)} of audio in ${wall.toFixed(1)} s ` +
    `(${(trace.length * session.dt / wall).toFixed(2)}× real time) -> ` +
    `${out} (${WAV_RATE / 1000} kHz, ${(trace.length * session.dt).toFixed(2)} s)`,
);
session.destroy();
process.exit(0);