#!/usr/bin/env node
/**
 * Grade model outputs against a written rubric using an OpenAI model as judge.
 *
 * Usage:
 *   node scripts/llm-judge.mjs OUTPUTS.jsonl RUBRIC.md [--min-score 4] [--out graded.jsonl] [--limit N]
 *
 *   OUTPUTS.jsonl  One JSON object per line:
 *                    {"id": "case-1", "input": "what the user asked",
 *                     "output": "what your system answered",
 *                     "reference": "optional known-good answer or facts"}
 *   RUBRIC.md      Plain-language grading criteria and what each score 1-5 means.
 *   --min-score    Lowest score that counts as a pass (default 4).
 *   --out          Write one graded JSON object per line to this file.
 *   --limit        Only grade the first N cases (handy while tuning the rubric).
 *
 * Environment variables:
 *   OPENAI_API_KEY      required, read by the openai SDK. Never hardcode it.
 *   OPENAI_JUDGE_MODEL  required. A current OpenAI model ID that supports JSON mode
 *                       (response_format json_object). No default, so changes are explicit.
 *
 * Requirements: Node.js 18+ and the official SDK: npm install openai
 *
 * Caveats: an LLM judge is a noisy grader. Spot-check its scores against your own
 * on a sample before trusting a pass rate, keep the judge model fixed when comparing
 * runs, and do not use the same model as both the system under test and the judge
 * if you can avoid it.
 *
 * Exit codes:
 *   0  every case scored at or above --min-score
 *   1  one or more cases scored below --min-score
 *   2  usage or configuration error
 *   3  API error (auth, rate limit after retries, server error)
 *   4  the judge returned a reply that was not valid score JSON for some case
 */
import { readFileSync, writeFileSync } from "node:fs";

function fail(message, code = 2) {
  process.stderr.write(`error: ${message}\n`);
  process.exit(code);
}

function parseArgs(argv) {
  const opts = { positional: [], minScore: 4, out: null, limit: Infinity };
  for (let i = 0; i < argv.length; i += 1) {
    const arg = argv[i];
    const next = () => {
      i += 1;
      if (argv[i] === undefined) fail(`${arg} needs a value`);
      return argv[i];
    };
    if (arg === "--min-score") opts.minScore = Number(next());
    else if (arg === "--out") opts.out = next();
    else if (arg === "--limit") opts.limit = Number(next());
    else if (arg.startsWith("--")) fail(`unknown option ${arg}`);
    else opts.positional.push(arg);
  }
  if (opts.positional.length !== 2) fail("usage: llm-judge.mjs OUTPUTS.jsonl RUBRIC.md [--min-score N] [--out FILE] [--limit N]");
  if (!(opts.minScore >= 1 && opts.minScore <= 5)) fail("--min-score must be between 1 and 5");
  if (!(opts.limit > 0)) fail("--limit must be a positive number");
  return opts;
}

function loadCases(path) {
  let text;
  try {
    text = readFileSync(path, "utf8");
  } catch (err) {
    fail(`cannot read ${path}: ${err.code ?? err.message}`);
  }
  const cases = [];
  text.split(/\r?\n/).forEach((line, index) => {
    if (!line.trim()) return;
    let item;
    try {
      item = JSON.parse(line);
    } catch (err) {
      fail(`${path} line ${index + 1}: invalid JSON (${err.message})`);
    }
    if (!item.id || typeof item.output !== "string") fail(`${path} line ${index + 1}: needs "id" and string "output"`);
    cases.push(item);
  });
  if (!cases.length) fail(`${path} has no cases`);
  return cases;
}

const SYSTEM = `You are a strict, consistent grader. You grade one response at a time against the rubric below.
Judge only what is in the response. Do not reward length or confident tone. If a reference is given, treat it as ground truth.

<rubric>
{{RUBRIC}}
</rubric>

Reply with only a JSON object of this shape:
{"score": <integer 1-5>, "reasons": "<two or three sentences naming the specific strengths or defects that set the score>"}`;

function userMessage(item) {
  const parts = [`<user_input>\n${item.input ?? ""}\n</user_input>`, `<response>\n${item.output}\n</response>`];
  if (item.reference) parts.push(`<reference>\n${item.reference}\n</reference>`);
  parts.push("Grade the response. Return the JSON object only.");
  return parts.join("\n\n");
}

function parseVerdict(text) {
  try {
    const verdict = JSON.parse(text);
    const score = Number(verdict.score);
    if (Number.isInteger(score) && score >= 1 && score <= 5) return { score, reasons: String(verdict.reasons ?? "") };
  } catch {
    // fall through
  }
  return null;
}

const opts = parseArgs(process.argv.slice(2));
const model = (process.env.OPENAI_JUDGE_MODEL ?? "").trim();
if (!model) fail("set OPENAI_JUDGE_MODEL to the judge model ID");
if (!process.env.OPENAI_API_KEY) fail("OPENAI_API_KEY is not set");

let OpenAI;
try {
  ({ default: OpenAI } = await import("openai"));
} catch {
  fail("the openai package is not installed (npm install openai)");
}

const cases = loadCases(opts.positional[0]).slice(0, opts.limit);
let rubric;
try {
  rubric = readFileSync(opts.positional[1], "utf8");
} catch (err) {
  fail(`cannot read ${opts.positional[1]}: ${err.code ?? err.message}`);
}
const system = SYSTEM.replace("{{RUBRIC}}", rubric.trim());
const client = new OpenAI({ maxRetries: 3 });

const graded = [];
let below = 0;
let unparsable = 0;
for (const item of cases) {
  let reply;
  try {
    const completion = await client.chat.completions.create({
      model,
      response_format: { type: "json_object" },
      messages: [
        { role: "system", content: system },
        { role: "user", content: userMessage(item) },
      ],
    });
    reply = completion.choices[0]?.message?.content ?? "";
  } catch (err) {
    fail(`API call failed on case ${item.id}: ${err.message}`, 3);
  }
  const verdict = parseVerdict(reply);
  let record;
  if (!verdict) {
    unparsable += 1;
    record = { id: item.id, status: "judge_error", raw: reply };
    process.stdout.write(`ERR  ${item.id}: judge reply was not valid score JSON\n`);
  } else {
    const pass = verdict.score >= opts.minScore;
    if (!pass) below += 1;
    record = { id: item.id, status: pass ? "pass" : "fail", score: verdict.score, reasons: verdict.reasons };
    process.stdout.write(`${pass ? "PASS" : "FAIL"} ${item.id}: ${verdict.score}/5  ${verdict.reasons}\n`);
  }
  graded.push({ ...record, judge_model: model });
}

if (opts.out) writeFileSync(opts.out, graded.map((g) => JSON.stringify(g)).join("\n") + "\n");
const scored = graded.filter((g) => typeof g.score === "number");
const mean = scored.length ? scored.reduce((sum, g) => sum + g.score, 0) / scored.length : 0;
process.stdout.write(`\n${scored.length - below}/${graded.length} passed (min score ${opts.minScore}), mean ${mean.toFixed(2)}, judge ${model}\n`);
if (unparsable) process.exit(4);
process.exit(below ? 1 : 0);