// The JSX judge: erase-mode equivalence with ts-blank-space (plus the JSX-type-argument
// deviation) and downlevel-mode equivalence with ts.transpileModule's classic runtime.
// See docs/correctness.md and the design note this was built from.
// Usage: node test/jsx.mjs [file.jsonl]
//        node test/jsx.mjs --dir <dir> [--dir <dir>...]
//        node test/jsx.mjs --self-test
import fs from "node:fs";
import path from "node:path";
import ts from "typescript";
import { blankSourceFile } from "../vendor/ts-blank-space/out/index.js";
import { transformSync, engine } from "../src/index.js";

const SK = ts.SyntaxKind;
const ALLOWED = new Set([" ", ";", "(", ")"]);
const root = path.resolve(import.meta.dirname, "..");

// ts-blank-space leaves a JSX tag's own type arguments in place; we blank them
// too, since that's the one deliberate deviation this judge enforces.
function blankJsxTypeArgs(sf, src, blanked) {
  const spans = [];
  (function visit(n) {
    if ((ts.isJsxOpeningElement(n) || ts.isJsxSelfClosingElement(n)) && n.typeArguments) {
      const start = n.typeArguments.pos - 1;
      const end = n.typeArguments.end + 1;
      if (src[start] !== "<" || src[end - 1] !== ">") throw new Error(`unexpected chars around JSX type arguments at ${start}`);
      spans.push([start, end]);
    }
    ts.forEachChild(n, visit);
  })(sf);
  if (spans.length === 0) return blanked;
  const chars = blanked.split("");
  for (const [start, end] of spans) {
    for (let i = start; i < end; i++) {
      const c = src[i];
      if (c !== "\n" && c !== "\r") chars[i] = " ";
    }
  }
  return chars.join("");
}

function oracleErase(sf, src) {
  const errors = [];
  const blanked = blankSourceFile(sf, (n) => errors.push(n.getStart(sf)));
  return { out: blankJsxTypeArgs(sf, src, blanked), errors };
}

function tsOutputForJsx(src, pragma, pragmaFrag) {
  return ts.transpileModule(src, {
    fileName: "case.jsx",
    compilerOptions: {
      jsx: ts.JsxEmit.React,
      jsxFactory: pragma,
      jsxFragmentFactory: pragmaFrag,
      target: ts.ScriptTarget.ESNext,
      module: ts.ModuleKind.ESNext,
      // Keeps every import as written and adds no "use strict", like the engine.
      verbatimModuleSyntax: true,
      alwaysStrict: false,
    },
  }).outputText;
}

function positional(src, out) {
  if (out.length !== src.length) return `length ${out.length} != input ${src.length}`;
  for (let i = 0; i < src.length; i++) {
    const a = src[i], b = out[i];
    if (a === b) continue;
    if (a === "\n" || a === "\r") return `line terminator changed at ${i}`;
    if (!ALLOWED.has(b)) return `invented ${JSON.stringify(b)} at ${i}`;
  }
  return null;
}

function splitLines(text) {
  return text.split(/\r\n|\r|\n/);
}

// erase.mjs's visitor, extended to cover JsxText content (tag/attribute names are already Identifiers).
function jsxJsAst(code) {
  const sf = ts.createSourceFile("out.jsx", code, ts.ScriptTarget.ESNext, false, ts.ScriptKind.JSX);
  const out = [];
  (function visit(n) {
    if (n.kind === SK.EmptyStatement || n.kind === SK.SemicolonClassElement) return;
    out.push(n.kind);
    if (ts.isIdentifier(n) || ts.isPrivateIdentifier(n) || ts.isLiteralExpression(n) || ts.isTemplateLiteralToken?.(n)) out.push(JSON.stringify(n.text));
    else if (ts.isJsxText(n)) out.push(JSON.stringify(n.text));
    ts.forEachChild(n, visit);
    out.push(-1);
  })(sf);
  return { ok: sf.parseDiagnostics.length === 0, ast: out.join(",") };
}

// Canonical plain-JS AST for downlevel mode, folding shapes TypeScript's emit
// and ours are allowed to differ on (string concatenation, inlined object
// spreads, property-name spelling) without that being a real divergence.
function isStrCanon(c) {
  return Array.isArray(c) && c[0] === "StrLit";
}

function propNameCanon(name) {
  if (ts.isIdentifier(name)) return ["PropName", name.text];
  if (ts.isStringLiteral(name)) return ["PropName", name.text];
  return null;
}

function hasProtoKey(objectLiteral) {
  return objectLiteral.properties.some((p) => p.name && (ts.isIdentifier(p.name) || ts.isStringLiteral(p.name)) && p.name.text === "__proto__");
}

function objectLiteralPropsCanon(objectLiteral) {
  const out = [];
  for (const p of objectLiteral.properties) {
    const spread = ts.isSpreadAssignment(p) && ts.skipParentheses(p.expression);
    if (spread && ts.isObjectLiteralExpression(spread) && !hasProtoKey(spread)) {
      out.push(...objectLiteralPropsCanon(spread));
    } else {
      out.push(toCanon(p));
    }
  }
  return out;
}

function toCanon(node) {
  if (node.kind === SK.EmptyStatement) return null;
  if (ts.isParenthesizedExpression(node)) return toCanon(node.expression);
  if (ts.isExportDeclaration(node) && !node.exportClause && !node.moduleSpecifier) return null;
  if (ts.isIdentifier(node) || ts.isPrivateIdentifier(node)) return [node.kind, node.text];
  if (ts.isStringLiteral(node) || node.kind === SK.NoSubstitutionTemplateLiteral) return ["StrLit", node.text];
  if (ts.isNumericLiteral(node)) return [node.kind, String(Number(node.text))];
  if (ts.isBinaryExpression(node) && node.operatorToken.kind === SK.PlusToken) {
    const left = toCanon(node.left), right = toCanon(node.right);
    if (isStrCanon(left) && isStrCanon(right)) return ["StrLit", left[1] + right[1]];
    return [node.kind, left, right];
  }
  if (ts.isObjectLiteralExpression(node)) return [node.kind, objectLiteralPropsCanon(node).filter((v) => v !== null)];
  if (ts.isPropertyAssignment(node)) return [node.kind, propNameCanon(node.name) ?? toCanon(node.name), toCanon(node.initializer)];
  if (ts.isShorthandPropertyAssignment(node)) return [node.kind, propNameCanon(node.name) ?? toCanon(node.name)];
  const kids = [];
  ts.forEachChild(node, (c) => {
    const v = toCanon(c);
    if (v !== null) kids.push(v);
  });
  return [node.kind, ts.isLiteralExpression(node) || ts.isTemplateLiteralToken(node) ? node.text : null, kids];
}

function canonicalJs(code) {
  const sf = ts.createSourceFile("out.js", code, ts.ScriptTarget.ESNext, false, ts.ScriptKind.JS);
  return { ok: sf.parseDiagnostics.length === 0, ast: JSON.stringify(toCanon(sf)) };
}

// Pure functions over already-computed candidate outputs, mirroring
// downlevel.mjs's judgeCase, so --self-test can feed fabricated candidates with no engine involved.
function judgeErase(src, wantErase, wantEraseErrors, candidate) {
  if (candidate.threw) return [`erase threw: ${candidate.threw.message}`];
  const { out: got, errors: gotErrors } = candidate;
  const posProblem = positional(src, got);
  if (posProblem) return [`erase positional: ${posProblem}`];
  const want = jsxJsAst(wantErase);
  let gotAst;
  try { gotAst = jsxJsAst(got); } catch (e) { return [`erase: TypeScript could not parse our output: ${e.message}`]; }
  if (want.ok && !gotAst.ok) return ["erase: JS parse failed where oracle output parses"];
  if (want.ast !== gotAst.ast) return ["erase: JSX AST differs from oracle output"];
  const wantStarts = wantEraseErrors.join();
  const gotStarts = gotErrors.map((e) => e.start).join();
  if (wantStarts !== gotStarts) return [`erase: errors want [${wantEraseErrors}] got [${gotErrors.map((e) => e.start)}]`];
  return [];
}

function judgeDownlevel(src, wantErase, pragma, pragmaFrag, candidate) {
  if (candidate.threw) return [`downlevel threw: ${candidate.threw.message}`];
  const { out: got, errors: gotErrors } = candidate;
  const srcLines = splitLines(src), gotLines = splitLines(got);
  if (gotLines.length !== srcLines.length) return [`downlevel line count: got ${gotLines.length}, input has ${srcLines.length}`];
  const gotAst = canonicalJs(got);
  if (!gotAst.ok) return ["downlevel: candidate does not parse as plain JS"];
  const want = tsOutputForJsx(wantErase, pragma, pragmaFrag);
  const wantAst = canonicalJs(want);
  if (wantAst.ast !== gotAst.ast) return ["downlevel: AST differs from ts.transpileModule reference"];
  if (gotErrors.length) return [`downlevel: unexpected errors ${JSON.stringify(gotErrors)}`];
  return [];
}

// Used when ts-blank-space itself reported an error on the input: no reference
// transpile is trustworthy, so only the "doesn't throw, preserves lines" half holds.
function judgeDownlevelLenient(src, candidate) {
  if (candidate.threw) return [`downlevel threw: ${candidate.threw.message}`];
  const srcLines = splitLines(src), gotLines = splitLines(candidate.out);
  if (gotLines.length !== srcLines.length) return [`downlevel line count (lenient): got ${gotLines.length}, input has ${srcLines.length}`];
  return [];
}

// Random chunking through transformSync on an iterable, same shape as erase.mjs's chunked().
function* parts(src) {
  let seed = src.length;
  for (let i = 0; i < src.length; ) {
    seed = (seed * 1103515245 + 12345) & 0x7fffffff;
    const n = 1 + (seed % 4096);
    yield src.slice(i, i + n);
    i += n;
  }
}

function chunked(src, opts) {
  let out = "";
  for (const piece of transformSync(parts(src), { ...opts, onError() {} })) out += piece;
  return out;
}

function runEngine(src, opts) {
  const errors = [];
  try {
    const out = transformSync(src, { ...opts, onError: (e) => errors.push({ start: e.start, end: e.end }) });
    return { out, errors };
  } catch (e) {
    return { threw: e };
  }
}

// Runs the real engine; --self-test never calls this, it calls judgeErase/judgeDownlevel(Lenient) directly.
function judgeCase(name, src, pragma, pragmaFrag) {
  const sf = ts.createSourceFile("case.tsx", src, ts.ScriptTarget.ESNext, true, ts.ScriptKind.TSX);
  if (sf.parseDiagnostics.length) return null; // skip: doesn't parse as TSX
  const { out: wantErase, errors: wantEraseErrors } = oracleErase(sf, src);

  const eraseCandidate = runEngine(src, { mode: "erase", jsx: true });
  const eraseProblems = judgeErase(src, wantErase, wantEraseErrors, eraseCandidate);
  if (eraseProblems.length === 0 && engine().push) {
    let streamed;
    try { streamed = chunked(src, { mode: "erase", jsx: true }); } catch (e) { eraseProblems.push(`erase streamed input threw: ${e.message}`); }
    if (streamed !== undefined && streamed !== eraseCandidate.out) eraseProblems.push("erase streamed output differs from whole-input output");
  }

  const downlevelCandidate = runEngine(src, { mode: "downlevel", jsx: true, pragma, pragmaFrag });
  // The engine only writes the classic runtime; these pragmas switch TypeScript to the automatic one.
  const automatic = ts.getLeadingCommentRanges(src, 0)?.some((r) => /@jsx(importsource|runtime)\b/i.test(src.slice(r.pos, r.end)));
  const downlevelProblems = wantEraseErrors.length === 0 && !automatic
    ? judgeDownlevel(src, wantErase, pragma, pragmaFrag, downlevelCandidate)
    : judgeDownlevelLenient(src, downlevelCandidate);
  if (downlevelProblems.length === 0 && engine().push) {
    let streamed;
    try { streamed = chunked(src, { mode: "downlevel", jsx: true, pragma, pragmaFrag }); } catch (e) { downlevelProblems.push(`downlevel streamed input threw: ${e.message}`); }
    if (streamed !== undefined && streamed !== downlevelCandidate.out) downlevelProblems.push("downlevel streamed output differs from whole-input output");
  }

  return [...eraseProblems, ...downlevelProblems];
}

function splitFilenameMarkers(src) {
  const lines = src.split(/\r\n|\r|\n/);
  const markers = [];
  lines.forEach((line, i) => {
    const m = /^\/\/\s*@filename\s*:\s*(.+?)\s*$/i.exec(line.trim());
    if (m) markers.push({ line: i, name: m[1] });
  });
  if (markers.length === 0) return null;
  return markers.map((marker, k) => {
    const start = marker.line + 1;
    const end = k + 1 < markers.length ? markers[k + 1].line : lines.length;
    return { filename: marker.name, src: lines.slice(start, end).join("\n") };
  });
}

function* walk(dir, top = true) {
  if (!top && fs.existsSync(path.join(dir, ".sweep-skip"))) return;
  let entries;
  try { entries = fs.readdirSync(dir, { withFileTypes: true }); } catch { return; }
  for (const e of entries) {
    if (e.name.startsWith(".") || e.name === "node_modules") continue;
    const full = path.join(dir, e.name);
    if (e.isDirectory()) { yield* walk(full, false); continue; }
    if (!e.name.endsWith(".tsx") || e.name.endsWith(".d.tsx")) continue;
    let src;
    try { src = fs.readFileSync(full, "utf8"); } catch { continue; }
    const segments = splitFilenameMarkers(src);
    if (!segments) { yield { name: full, src }; continue; }
    for (const seg of segments) {
      if (!seg.filename.endsWith(".tsx") || seg.filename.endsWith(".d.tsx")) continue;
      yield { name: `${full}::${seg.filename}`, src: seg.src };
    }
  }
}

function* inputs(argv) {
  if (argv[0] === "--dir") {
    for (let i = 0; i < argv.length; i++) if (argv[i] === "--dir") yield* walk(path.resolve(argv[++i]));
    return;
  }
  const file = argv[0] ?? "corpus/jsx.jsonl";
  const lines = fs.readFileSync(file, "utf8").split("\n").filter(Boolean);
  for (const [i, line] of lines.entries()) {
    const c = JSON.parse(line);
    yield { name: c.name ?? `${path.basename(file)}/${i}`, src: c.src, pragma: c.pragma, pragmaFrag: c.pragmaFrag };
  }
}

function run(argv) {
  let pass = 0, skipped = 0;
  const failures = [];
  const tally = {};
  for (const { name, src, pragma, pragmaFrag } of inputs(argv)) {
    const problems = judgeCase(name, src, pragma, pragmaFrag);
    if (problems === null) { skipped++; continue; }
    if (problems.length === 0) { pass++; continue; }
    const kind = problems[0].split(" ").slice(0, 3).join(" ");
    tally[kind] = (tally[kind] ?? 0) + 1;
    failures.push([name, problems.join("; ")]);
  }

  const baselinePath = process.env.JSX_BASELINE;
  if (baselinePath && process.env.JSX_BASELINE_WRITE) fs.writeFileSync(baselinePath, failures.map(([n]) => path.relative(process.cwd(), n)).sort().join("\n") + "\n");
  const known = new Set(baselinePath && fs.existsSync(baselinePath) ? fs.readFileSync(baselinePath, "utf8").split("\n").filter(Boolean) : []);
  const isKnown = ([n]) => known.has(path.relative(process.cwd(), n));
  const fresh = failures.filter((f) => !isKnown(f));
  for (const [n, p] of fresh.slice(0, Number(process.env.SHOW ?? 15))) console.log(`FAIL ${n}: ${p}`);
  console.log(`jsx: ${pass} pass, ${failures.length} fail, ${skipped} skipped`, failures.length ? JSON.stringify(tally) : "");
  if (known.size) {
    const failing = new Set(failures.map(([n]) => path.relative(process.cwd(), n)));
    const fixed = [...known].filter((n) => !failing.has(n));
    console.log(`baseline: ${known.size} known, ${fresh.length} new, ${fixed.length} now passing${fixed.length ? " (shrink the baseline)" : ""}`);
  }
  process.exitCode = fresh.length ? 1 : 0;
}

function fixtureFor(src) {
  const sf = ts.createSourceFile("case.tsx", src, ts.ScriptTarget.ESNext, true, ts.ScriptKind.TSX);
  const { out: wantErase, errors: wantEraseErrors } = oracleErase(sf, src);
  return { sf, wantErase, wantEraseErrors };
}

function selfTest() {
  let ok = true;
  const check = (label, pass) => {
    console.log(`${pass ? "ok" : "FAIL"} ${label}`);
    ok = ok && pass;
    return pass;
  };

  const eraseFixtures = [
    "export const a = <div a=\"x\">hi</div>;\n",
    "export const b = <Foo<string> x=\"1\">hi</Foo>;\n",
    "export const c = <>{a}<b/></>;\n",
  ];
  for (const src of eraseFixtures) {
    const { wantErase, wantEraseErrors } = fixtureFor(src);
    const problems = judgeErase(src, wantErase, wantEraseErrors, { out: wantErase, errors: wantEraseErrors.map((start) => ({ start })) });
    check(`erase accepts the oracle's own output (${JSON.stringify(src)})`, problems.length === 0);
  }

  {
    const src = "export const b = <Foo<string> x=\"1\">hi</Foo>;\n";
    const { sf, wantErase, wantEraseErrors } = fixtureFor(src);
    const oracleRawOutput = blankSourceFile(sf, () => {});
    const problems = judgeErase(src, wantErase, wantEraseErrors, { out: oracleRawOutput, errors: [] });
    check("erase rejects a candidate that left JSX type arguments unblanked", problems.length > 0);
  }

  const downlevelFixtures = [
    "export const a = <div a=\"x\">hi</div>;\n",
    "export const b = <>{a}</>;\n",
    "export const c = <div {...{x: 1}} y>{(1, 2)}</div>;\n",
  ];
  for (const src of downlevelFixtures) {
    const { wantErase } = fixtureFor(src);
    const want = tsOutputForJsx(wantErase, undefined, undefined);
    const problems = judgeDownlevel(src, wantErase, undefined, undefined, { out: want, errors: [] });
    check(`downlevel accepts ts.transpileModule's own output (${JSON.stringify(src)})`, problems.length === 0);
  }

  {
    const src = "export const a = <div>x</div>;\nexport const b = <div>hi there</div>;\n";
    const { wantErase } = fixtureFor(src);
    const correct = tsOutputForJsx(wantErase, undefined, undefined);

    const sanity = judgeDownlevel(src, wantErase, undefined, undefined, { out: correct, errors: [] });
    check("downlevel accepts the correct two-line candidate", sanity.length === 0);

    const broken = {
      "wrong factory name": correct.replace("React.createElement(\"div\", null, \"x\")", "Foo.createElement(\"div\", null, \"x\")"),
      "dropped child": correct.replace("React.createElement(\"div\", null, \"x\")", "React.createElement(\"div\", null)"),
      "doubled whitespace inside JSX text": correct.replace("hi there", "hi  there"),
      "a line removed": correct.replace("\n", " "),
    };
    for (const [label, out] of Object.entries(broken)) {
      const problems = judgeDownlevel(src, wantErase, undefined, undefined, { out, errors: [] });
      check(`downlevel rejects "${label}"`, problems.length > 0);
    }
  }

  console.log(ok ? "self-test: pass" : "self-test: FAIL");
  process.exitCode = ok ? 0 : 1;
}

const argv = process.argv.slice(2);
if (argv[0] === "--self-test") selfTest();
else run(argv);