/**
 * ECSeg browser benchmark driver — measures ECSeg ONNX variants in the TARGET runtime
 * (onnxruntime-web / WASM CPU EP), the only measurement that decides a production swap. Python
 * onnxruntime latency is explicitly NOT a substitute (WASM has different kernels/threading).
 *
 * It serves the harness + ORT dist + models + a precomputed input tensor to a real browser and
 * drives scripts/model-optimization/bench/harness.mjs, which configures ORT exactly as the app's
 * ECSeg session does (executionProviders ['cpu'], graphOptimizationLevel 'all', logSeverityLevel 3,
 * executionMode sequential|parallel by thread count, simd on).
 *
 * Two configurations mirror the two ways a real ECSeg deployment is threaded:
 *   wasm-mt : cross-origin isolated (COOP/COEP) => SharedArrayBuffer => ORT auto-sizes a worker
 *             pool. This is app.annotateit.ai in a browser.
 *   wasm-st : not isolated => one thread. This is the desktop/iOS build and any origin without
 *             COOP/COEP.
 *
 * Usage:
 *   node bench_browser.mjs --models ecseg-s.fp32.onnx,ecseg-s.fp16.onnx --configs wasm-mt,wasm-st \
 *        --engine chrome --repeats 12 --out results-s.json
 */

import { execSync } from 'node:child_process';
import fs from 'node:fs';
import http from 'node:http';
import os from 'node:os';
import path from 'node:path';
import { fileURLToPath } from 'node:url';

import { chromium, webkit } from 'playwright';

const HERE = path.dirname(fileURLToPath(import.meta.url));
const ROOT = path.resolve(HERE, '..', '..', '..');
const ORT_DIST = path.join(ROOT, 'node_modules', 'onnxruntime-web', 'dist');

const arg = (name, fallback) => {
    const i = process.argv.indexOf(`--${name}`);
    return i === -1 ? fallback : process.argv[i + 1];
};

const MODELS_DIR = arg('models-dir', '');
const INPUT_DIR = arg('input-dir', '');
const MODELS = (arg('models', '') ?? '').split(',').filter(Boolean);
const CONFIG_IDS = (arg('configs', 'wasm-mt,wasm-st') ?? '').split(',').filter(Boolean);
const ENGINE = arg('engine', 'chrome');
const REPEATS = Number(arg('repeats', 12));
const WARMUPS = Number(arg('warmups', 2));
const INPUT = arg('input', '000000000139.f32');
const OUT_FILE = arg('out', path.join(HERE, 'results.json'));
const CELL_TIMEOUT_MS = Number(arg('cellTimeoutMs', 300_000));
// Diagnostic overrides: force a thread count / executionMode regardless of the config's defaults.
const THREADS_OVERRIDE = arg('threads', '');
const EM_OVERRIDE = arg('em', '');

if (!MODELS_DIR || !INPUT_DIR) throw new Error('pass --models-dir and --input-dir');

const ALL_CONFIGS = {
    'wasm-mt': { id: 'wasm-mt', threads: 'auto', isolated: true },
    'wasm-st': { id: 'wasm-st', threads: '1', isolated: false },
};
const CONFIGS = CONFIG_IDS.map((id) => {
    if (!ALL_CONFIGS[id]) throw new Error(`unknown config ${id}`);
    return ALL_CONFIGS[id];
});

const CONTENT_TYPES = {
    '.html': 'text/html; charset=utf-8',
    '.mjs': 'text/javascript; charset=utf-8',
    '.js': 'text/javascript; charset=utf-8',
    '.wasm': 'application/wasm',
    '.map': 'application/json',
    '.onnx': 'application/octet-stream',
    '.f32': 'application/octet-stream',
};

const startServer = (isolated) =>
    new Promise((resolve) => {
        const server = http.createServer((req, res) => {
            const url = new URL(req.url, 'http://localhost');
            const pathname = decodeURIComponent(url.pathname);

            let file;
            if (pathname === '/' || pathname === '/harness.html') file = path.join(HERE, 'harness.html');
            else if (pathname === '/harness.mjs') file = path.join(HERE, 'harness.mjs');
            else if (pathname.startsWith('/ort/')) file = path.join(ORT_DIST, pathname.slice(5));
            else if (pathname.startsWith('/models/')) file = path.join(MODELS_DIR, pathname.slice(8));
            else if (pathname.startsWith('/input/')) file = path.join(INPUT_DIR, pathname.slice(7));
            else file = '';

            if (!file || !fs.existsSync(file) || fs.statSync(file).isDirectory()) {
                if (process.env.BENCH_DEBUG) console.error(`    [server 404] ${pathname}`);
                res.writeHead(404).end('not found');
                return;
            }
            const headers = { 'Content-Type': CONTENT_TYPES[path.extname(file)] ?? 'application/octet-stream' };
            if (isolated) {
                headers['Cross-Origin-Opener-Policy'] = 'same-origin';
                headers['Cross-Origin-Embedder-Policy'] = 'require-corp';
                headers['Cross-Origin-Resource-Policy'] = 'same-origin';
            }
            res.writeHead(200, headers);
            fs.createReadStream(file).pipe(res);
        });
        server.listen(0, '127.0.0.1', () => resolve({ server, port: server.address().port }));
    });

// Bringing up ORT's threaded WASM worker pool is finicky; the proven scripts/benchmark-sam2.mjs
// recipe is a HEADED persistent context. We mirror it exactly for Chrome (a persistent user-data-dir
// + headed is what lets the pool start on this machine). BENCH_HEADLESS=1 forces headless (fine for
// the single-thread configs, which never spawn a pool).
const HEADLESS = process.env.BENCH_HEADLESS === '1';
let udidCounter = 0;
// Returns { browser } where browser has .newPage() and .close(), regardless of launch mode.
const launch = async () => {
    if (ENGINE === 'webkit') return webkit.launch({ headless: HEADLESS });
    const userDataDir = path.join(os.tmpdir(), `ecseg-bench-${process.pid}-${udidCounter++}`);
    fs.rmSync(userDataDir, { recursive: true, force: true });
    return chromium.launchPersistentContext(userDataDir, {
        channel: 'chrome',
        headless: HEADLESS,
        args: ['--force-high-performance-gpu', '--no-first-run', '--no-default-browser-check'],
    });
};

const hostInfo = () => {
    let cpu = os.cpus()?.[0]?.model ?? 'unknown';
    let chip = '';
    try { chip = execSync('sysctl -n machdep.cpu.brand_string', { encoding: 'utf8' }).trim(); } catch {}
    return {
        platform: `${os.platform()} ${os.release()}`,
        arch: os.arch(),
        cpu: chip || cpu,
        logicalCores: os.cpus()?.length ?? null,
        totalMemGB: Math.round(os.totalmem() / 2 ** 30),
    };
};

const runCell = async (browser, port, model, config) => {
    const page = await browser.newPage();
    if (process.env.BENCH_DEBUG) {
        page.on('console', (m) => console.error(`    [page:${m.type()}] ${m.text()}`));
        page.on('pageerror', (e) => console.error(`    [pageerror] ${e.message}`));
        page.on('requestfailed', (r) => console.error(`    [reqfail] ${r.url()} ${r.failure()?.errorText}`));
        page.on('response', (r) => { if (r.status() >= 400) console.error(`    [http ${r.status()}] ${r.url()}`); });
    }
    try {
        const threads = THREADS_OVERRIDE || config.threads;
        const url =
            `http://127.0.0.1:${port}/harness.html` +
            `?model=${encodeURIComponent(model)}&input=${encodeURIComponent(INPUT)}` +
            `&threads=${threads}&repeats=${REPEATS}&warmups=${WARMUPS}` +
            (EM_OVERRIDE ? `&em=${EM_OVERRIDE}` : '');
        await page.goto(url, { waitUntil: 'load' });

        const payload = await Promise.race([
            page.evaluate(() => globalThis.benchmarkPromise),
            new Promise((resolve) =>
                setTimeout(
                    () =>
                        page
                            .evaluate(() => document.getElementById('log')?.textContent ?? 'unknown')
                            .catch(() => 'unknown')
                            .then((phase) => resolve({ ok: false, timedOut: true, stalledAt: phase })),
                    CELL_TIMEOUT_MS
                )
            ),
        ]);
        return { model, config: config.id, ...payload };
    } finally {
        await page.close();
    }
};

const main = async () => {
    const host = { ...hostInfo(), engine: ENGINE, ortWeb: '1.24.3', when: new Date().toISOString() };
    console.log(`ECSeg browser benchmark — ${ENGINE} — ${host.cpu} (${host.logicalCores} cores)`);

    const cells = [];
    for (const config of CONFIGS) {
        const { server, port } = await startServer(config.isolated);
        const browser = await launch();
        try {
            for (const model of MODELS) {
                process.stdout.write(`  ${model.padEnd(34)} ${config.id.padEnd(8)} … `);
                try {
                    const cell = await runCell(browser, port, model, config);
                    if (cell.ok === false) {
                        console.log(cell.timedOut ? `STALLED @ "${cell.stalledAt}"` : `FAILED: ${String(cell.error).split('\n')[0]}`);
                    } else {
                        const r = cell.result;
                        console.log(
                            `create ${Math.round(r.sessionCreateMs)}ms · cold ${Math.round(r.coldInferenceMs)}ms · ` +
                            `warm p50 ${Math.round(r.warm.p50)}ms (p90 ${Math.round(r.warm.p90)}) · inst ${r.fingerprint.numInstances}` +
                            (r.fingerprint.anyNaNInf ? ' · NaN!' : '')
                        );
                    }
                    cells.push(cell);
                } catch (e) {
                    console.log(`error: ${e.message}`);
                    cells.push({ model, config: config.id, ok: false, error: e.message });
                }
            }
        } finally {
            await browser.close();
            server.close();
        }
    }

    fs.writeFileSync(OUT_FILE, JSON.stringify({ host, repeats: REPEATS, warmups: WARMUPS, input: INPUT, cells }, null, 2));
    console.log(`\n  wrote ${OUT_FILE}`);
};

await main();