// Run the sentence helper on the eval exchanges (items.json) in a real
// browser on a real graphics chip, with the app's own prompts and runtime,
// and write every answer for reading.
//
//   PLAYWRIGHT=<path to playwright's index.mjs> node tools/ai/eval/run.mjs SITE OUT.json [lang...]
//
// SITE is a built site holding vendor/litert-lm-*/ and ai/ (for example
// aac-aurasense-web/dist/app, or the directory prerender-static.sh builds).
// Languages default to every language with examples, ready or not: from the
// catalogue, or from ../languages.json with LANGUAGES=1 (to test languages
// before they are published). ITEMS=<file> reads other exchanges than
// items.json. The browser is installed Chrome with WebGPU (ENGINE=webkit for
// Playwright's WebKit, whose 1 GB storage limit is too small for the model).
// MODEL=<id> chooses the model where the site has more than one.
//
// The prompts are ai_helper::prompt's, word for word (a test in ai_helper.rs
// checks the instructions below against it). Answers are greedy, so a rerun
// with the same model, runtime and prompts gives the same answers. Reading
// them is the test: add "usable": true or false to each sentence answer, as in
// results-2026-10-01.json, before marking a language ready in
// ../languages.json.
import fs from 'node:fs';
import http from 'node:http';
import path from 'node:path';
import { fileURLToPath } from 'node:url';

const HERE = path.dirname(fileURLToPath(import.meta.url));
const BINDINGS = path.resolve(HERE, '../../../crates/aac-web/src/bindings');
const [siteArg, outArg, ...langArgs] = process.argv.slice(2);
if (!siteArg || !outArg) {
  console.error('usage: node tools/ai/eval/run.mjs SITE OUT.json [lang...]');
  process.exit(2);
}
const SITE = path.resolve(siteArg);
const { chromium, webkit } = await import(process.env.PLAYWRIGHT || 'playwright');

// ai_helper::prompt's instructions, word for word.
const SENTENCE_SYSTEM = 'You help a person with a speech disability talk through an AAC app. They choose a few key words, in {lang}. Write the full sentence they most likely want to say, fitting what their partner said when a partner line is given. Answer with 3 different sentences in {lang}, one per line, most likely first, and nothing else.';
const REPLY_SYSTEM = 'You help a person with a speech disability talk through an AAC app. Their partner just spoke to them, in {lang}. Suggest 4 short, different things the person might say back, in the first person, in {lang}: cover different answers (for example yes and no, or different choices), most useful first, one per line, and nothing else.';
const MAX_OUTPUT_TOKENS = 128;

const line = (partner, words) =>
  partner && words ? `Partner: ${partner}\nWords: ${words}` : words ? `Words: ${words}` : `Partner: ${partner}`;

function prompt(kind, lang, partner, words) {
  const preface = [{ role: 'system', content: (kind === 'reply' ? REPLY_SYSTEM : SENTENCE_SYSTEM).replaceAll('{lang}', lang.name) }];
  if (kind === 'reply') {
    for (const e of lang.examples.reply) preface.push({ role: 'user', content: line(e.partner, null) }, { role: 'assistant', content: e.out.join('\n') });
  } else {
    for (const e of lang.examples.sentence) preface.push({ role: 'user', content: line(e.partner, e.words) }, { role: 'assistant', content: e.out.join('\n') });
  }
  return { preface, user: line(partner, words), max_output_tokens: MAX_OUTPUT_TOKENS };
}

const pointer = JSON.parse(fs.readFileSync(path.join(SITE, 'ai/index.json'), 'utf8'));
const catalogue = JSON.parse(fs.readFileSync(path.join(SITE, 'ai', pointer.current), 'utf8'));
if (process.env.LANGUAGES) {
  catalogue.languages = JSON.parse(fs.readFileSync(path.join(HERE, '../languages.json'), 'utf8')).languages;
}
const items = JSON.parse(fs.readFileSync(process.env.ITEMS || path.join(HERE, 'items.json'), 'utf8'));
const langs = langArgs.length ? langArgs : Object.keys(catalogue.languages).filter((c) => catalogue.languages[c].examples?.sentence?.length && items[c]);
// One model per run: MODEL=<id> when the site has more than one (each
// language is tested on every model that might serve it).
const ids = Object.keys(catalogue.models);
const modelId = process.env.MODEL || (ids.length === 1 ? ids[0] : '');
if (!catalogue.models[modelId]) {
  console.error(`set MODEL to one of: ${ids.join(', ')}`);
  process.exit(2);
}
const model = catalogue.models[modelId];
const vendor = fs.readdirSync(path.join(SITE, 'vendor')).find((d) => d.startsWith('litert-lm-'));

const tasks = [];
for (const code of langs) {
  const lang = catalogue.languages[code];
  for (const [partner, target, words] of items[code] ?? []) {
    tasks.push({ kind: 'sentence', lang: code, partner, words, target, prompt: prompt('sentence', lang, partner, words) });
    tasks.push({ kind: 'sentence0', lang: code, partner: null, words, target, prompt: prompt('sentence', lang, null, words) });
    tasks.push({ kind: 'reply', lang: code, partner, words: null, target, prompt: prompt('reply', lang, partner, null) });
  }
}

const PAGE = `<!doctype html><meta charset="utf-8"><title>sentence helper eval</title>
<script type="module">
import * as ai from "/bindings/ai_helper.js";
const STEMS = { "jspi": "litertlm_wasm_internal", "asyncify": "litertlm_wasm_asyncify_internal",
                "compat-jspi": "litertlm_wasm_compat_internal", "compat-asyncify": "litertlm_wasm_compat_asyncify_internal" };
const hex = (b) => [...new Uint8Array(b)].map((x) => x.toString(16).padStart(2, "0")).join("");
async function store(url, sha256) {
  const cache = await caches.open("aac-models-v1");
  const key = new URL(url, location.href).href;
  if (await cache.match(key)) return;
  const bytes = await (await fetch(key, { cache: "no-store" })).arrayBuffer();
  if (sha256 && hex(await crypto.subtle.digest("SHA-256", bytes)) !== sha256) throw new Error("sha256 mismatch: " + url);
  await cache.put(key, new Response(bytes));
}
window.run = async ({ vendor, parts, key, tasks }) => {
  const f = JSON.parse(await ai.wasm_features());
  const build = f.relaxed_simd ? (f.jspi ? "jspi" : "asyncify") : (f.jspi ? "compat-jspi" : "compat-asyncify");
  const base = "/app/vendor/" + vendor;
  const files = { worker: base + "/ai-worker.js", glue: base + "/" + STEMS[build] + ".js", wasm: base + "/" + STEMS[build] + ".wasm" };
  for (const u of Object.values(files)) await store(u);
  for (const p of parts) await store("/app/ai/" + p.url, p.sha256);
  const urls = parts.map((p) => new URL("/app/ai/" + p.url, location.href).href);
  await ai.ensure_engine(base, JSON.stringify(files), key, JSON.stringify(urls), 1024);
  const out = [];
  for (const t of tasks) {
    const t0 = performance.now();
    const text = await ai.generate(JSON.stringify(t.prompt));
    out.push({ ms: Math.round(performance.now() - t0), text });
  }
  ai.unload();
  await caches.delete("aac-models-v1");
  return { build, out };
};
window.ready = true;
</script>`;

const TYPES = { '.js': 'text/javascript', '.wasm': 'application/wasm', '.json': 'application/json' };
const server = http.createServer((req, res) => {
  const url = decodeURIComponent(new URL(req.url, 'http://x').pathname);
  if (url === '/') { res.writeHead(200, { 'Content-Type': 'text/html' }); return res.end(PAGE); }
  const file = url.startsWith('/bindings/') ? path.join(BINDINGS, url.slice(10))
    : url.startsWith('/app/') ? path.join(SITE, url.slice(5)) : null;
  if (!file || !fs.existsSync(file) || !fs.statSync(file).isFile()) { res.writeHead(404); return res.end(); }
  res.writeHead(200, { 'Content-Type': TYPES[path.extname(file)] ?? 'application/octet-stream' });
  fs.createReadStream(file).pipe(res);
});
await new Promise((r) => server.listen(0, '127.0.0.1', r));
const origin = `http://127.0.0.1:${server.address().port}`;

const browser = process.env.ENGINE === 'webkit'
  ? await webkit.launch({ headless: true })
  : await chromium.launch({ channel: 'chrome', headless: true, args: ['--enable-unsafe-webgpu'] });
const page = await browser.newPage();
await page.goto(origin + '/');
await page.waitForFunction(() => window.ready === true);
const result = await page.evaluate((arg) => window.run(arg), {
  vendor, parts: model.parts, key: `${modelId}@${model.parts[0].sha256}`, tasks: tasks.map((t) => ({ prompt: t.prompt })),
});
await browser.close();
server.close();

const answers = tasks.map((t, i) => ({ kind: t.kind, lang: t.lang, partner: t.partner, words: t.words, target: t.target, ...result.out[i] }));
fs.writeFileSync(outArg, JSON.stringify({
  about: `The sentence helper's answers for items.json: ${model.name} on ${vendor} (${result.build}), the app's prompts, greedy, ${new Date().toISOString().slice(0, 10)}.`,
  answers,
}, null, 1) + '\n');
console.log(`✓ ${answers.length} answers in ${langs.join(', ')} → ${outArg}`);