apps/api/src/judge/runner.ts

import type { Config } from "../config.js";
import { AiClient, type AiClientConfig } from "../ai/client.js";
import type { OptionsSnapshot } from "../queue/runner.js";
import { DEFAULT_JUDGE_PROMPT_TEMPLATE, JSON_ONLY_SYSTEM } from "./prompt.js";
import { buildGoldenTrajectory } from "./trajectory.js";
import { evaluateJudgePolicy } from "./policy.js";

type AnyObj = Record<string, any>;

const DIMENSION_LABELS: Record<string, string> = {
  task_realism: "Task Realism",
  no_information_leakage: "No Information Leakage",
  task_complexity: "Task Complexity",
  rubric_concreteness: "Rubric Concreteness",
  rubric_completeness: "Rubric Completeness",
  trajectory_task_alignment: "Trajectory-Task Alignment",
};

const DIMENSION_ORDER = [
  "task_realism",
  "no_information_leakage",
  "task_complexity",
  "rubric_concreteness",
  "rubric_completeness",
  "trajectory_task_alignment",
] as const;

const SUPPORTED_VARIABLES = ["taskStatement", "rubricJsonStr", "goldenTrajectory", "passPolicy"];

export interface RunJudgeInput {
  aiClient: AiClient;
  config: Config;
  snapshot: OptionsSnapshot;
  taskPackage: AnyObj;
  evidence: AnyObj;
}

export interface JudgeResult {
  dimensions: Record<string, { score: number; explanation: string }>;
  total_score: number;
  max_score: number;
  has_zeros: boolean;
  trajectory_pass: boolean;
  dbdiff_pass: boolean;
  package_pass: boolean;
  verdict: "PASS" | "NEEDS_REVISION" | "REJECT";
  policy: {
    strictFullMarks: boolean;
    ignoreTaskComplexityForFullMarks: boolean;
    threshold: number;
    maxScore: number;
    rule: string;
  };
  summary: string;
  revision_notes: string[];
  dimensionLabels: Record<string, string>;
  raw?: unknown;
  evaluated?: boolean;
  evaluatedAt?: string;
  overallScore?: number;
  maxScore?: number;
  passed?: boolean;
  feedback?: string;
  rubricScores?: Array<{
    rubricId: string;
    name: string;
    score: number;
    maxScore: number;
    feedback: string;
  }>;
  revisionNotes?: string[];
  promptUsed?: string;
  rawResponse?: unknown;
}

type ParsedJudgeResponse = {
  dimensions: Record<(typeof DIMENSION_ORDER)[number], { score: 0 | 1 | 2; explanation: string }>;
  total_score: number;
  has_zeros: boolean;
  verdict: "PASS" | "NEEDS_REVISION" | "REJECT";
  summary: string;
  revision_notes: string[];
};

function getUnknownTemplateVariables(template: string): string[] {
  const matches = String(template || "").match(/\$\{([^}]+)\}/g) ?? [];
  return matches
    .map((match) => match.slice(2, -1).trim())
    .filter((name) => !SUPPORTED_VARIABLES.includes(name));
}

function buildJudgePrompt(template: string, variables: Record<string, string>): string {
  let prompt = String(template || "");
  for (const [key, value] of Object.entries(variables)) {
    const token = "${" + key + "}";
    const escapedToken = "\\" + token;
    prompt = prompt.split(escapedToken).join(value);
    prompt = prompt.split(token).join(value);
  }
  return prompt;
}

function buildPassPolicyText(snapshot: OptionsSnapshot): string {
  if (snapshot.strictFullMarks) {
    if (snapshot.ignoreTaskComplexityForFullMarks) {
      return [
        "Strict full-marks mode is enabled with Task Complexity ignored for full-marks evaluation.",
        "PASS if all dimensions except task_complexity score 2/2 and task_complexity scores at least 1/2.",
        "A task_complexity score of 1/2 is acceptable; a task_complexity score of 0/2 is not acceptable.",
        "Any non-task_complexity score below 2 must be NEEDS_REVISION or REJECT according to severity.",
      ].join("\n");
    }
    return [
      "Strict full-marks mode is enabled.",
      "PASS only if total_score equals 12 and has_zeros is false.",
      "Any score below 12 must be NEEDS_REVISION or REJECT according to severity.",
    ].join("\n");
  }
  return [
    "Normal mode is enabled.",
    "PASS if total_score >= 10 and has_zeros is false.",
    "NEEDS_REVISION if total_score is 7-9 or total_score >= 10 with has_zeros true.",
    "REJECT if total_score < 7.",
  ].join("\n");
}

function buildRubricJson(rubrics: AnyObj[]): string {
  return JSON.stringify(
    (rubrics || []).map((rubric) => ({
      id: rubric.id,
      name: rubric.name,
      description: rubric.description,
      category: rubric.category ?? "",
      checker_key: rubric.checker_key ?? "",
      maxScore: rubric.maxScore,
      scorePoints: rubric.scorePoints,
    })),
    null,
    2,
  );
}

function normalizeExistingGoldenTrajectory(value: unknown): unknown[] | string | null {
  if (Array.isArray(value)) return value.length ? value : null;
  if (typeof value !== "string") return null;
  const trimmed = value.trim();
  if (!trimmed || trimmed === "[]") return null;
  try {
    const parsed = JSON.parse(trimmed) as unknown;
    return Array.isArray(parsed) && parsed.length ? parsed : trimmed;
  } catch (_) {
    return trimmed;
  }
}

function resolveGoldenTrajectory(taskPackage: AnyObj, evidence: AnyObj, groups: AnyObj[]): string {
  const built = buildGoldenTrajectory(evidence, groups);
  if (built && built.trim() && built.trim() !== "[]") return built;
  if (evidence?.type === "mcp") return "[]";
  const existing = normalizeExistingGoldenTrajectory(taskPackage?.golden_trajectory);
  if (existing) {
    return typeof existing === "string" ? existing : JSON.stringify(existing, null, 2);
  }
  return "[]";
}

function stripCodeFence(content: string): string {
  return String(content || "")
    .replace(/^```json\s*/i, "")
    .replace(/^```\s*/i, "")
    .replace(/```\s*$/i, "");
}

function extractJsonObject(content: string): string {
  const start = content.indexOf("{");
  const end = content.lastIndexOf("}");
  if (start === -1 || end === -1 || end <= start) return "";
  return content.slice(start, end + 1);
}

function invalid(): never {
  throw new Error("Judge response JSON schema invalid");
}

function validateDimension(value: unknown): { score: 0 | 1 | 2; explanation: string } {
  if (!value || typeof value !== "object") invalid();
  const dimension = value as AnyObj;
  if (
    (dimension.score !== 0 && dimension.score !== 1 && dimension.score !== 2) ||
    typeof dimension.explanation !== "string"
  ) {
    invalid();
  }
  return {
    score: dimension.score,
    explanation: dimension.explanation,
  };
}

function validateTotalScore(value: unknown): number {
  if (typeof value !== "number" || value < 0 || value > 12) invalid();
  return value;
}

function validateVerdict(value: unknown): ParsedJudgeResponse["verdict"] {
  if (value === "PASS" || value === "NEEDS_REVISION" || value === "REJECT") return value;
  invalid();
}

function validateJudgeResponse(value: unknown): ParsedJudgeResponse {
  if (!value || typeof value !== "object") invalid();
  const response = value as AnyObj;
  const dimensions = response.dimensions;
  if (!dimensions || typeof dimensions !== "object") invalid();
  return {
    dimensions: {
      task_realism: validateDimension((dimensions as AnyObj).task_realism),
      no_information_leakage: validateDimension((dimensions as AnyObj).no_information_leakage),
      task_complexity: validateDimension((dimensions as AnyObj).task_complexity),
      rubric_concreteness: validateDimension((dimensions as AnyObj).rubric_concreteness),
      rubric_completeness: validateDimension((dimensions as AnyObj).rubric_completeness),
      trajectory_task_alignment: validateDimension((dimensions as AnyObj).trajectory_task_alignment),
    },
    total_score: validateTotalScore(response.total_score),
    has_zeros: typeof response.has_zeros === "boolean" ? response.has_zeros : invalid(),
    verdict: validateVerdict(response.verdict),
    summary: typeof response.summary === "string" ? response.summary : invalid(),
    revision_notes:
      Array.isArray(response.revision_notes) &&
      response.revision_notes.every((item) => typeof item === "string")
        ? response.revision_notes
        : invalid(),
  };
}

function parseJudgeResponse(content: string): ParsedJudgeResponse {
  const candidates = [
    content.trim(),
    stripCodeFence(content).trim(),
    extractJsonObject(content).trim(),
  ].filter(Boolean);
  for (const candidate of candidates) {
    try {
      return validateJudgeResponse(JSON.parse(candidate));
    } catch (_) {
      // try next candidate
    }
  }
  throw new Error("Judge response is not valid JSON");
}

function mapJudgeResponseToResult(
  response: ParsedJudgeResponse,
  promptUsed: string,
  rawResponse: unknown,
  snapshot: OptionsSnapshot,
): JudgeResult {
  const dimensions = response.dimensions;
  const rubricScores = DIMENSION_ORDER.map((key) => ({
    rubricId: "judge:" + key,
    name: DIMENSION_LABELS[key],
    score: dimensions[key].score,
    maxScore: 2,
    feedback: dimensions[key].explanation,
  }));
  const feedbackSections = [response.summary.trim()];
  if (response.revision_notes.length > 0) {
    feedbackSections.push(
      "Revision Notes:\n" + response.revision_notes.map((note) => "- " + note).join("\n"),
    );
  }
  const strictFullMarks = Boolean(snapshot.strictFullMarks);
  const policyDecision = evaluateJudgePolicy(
    {
      total_score: response.total_score,
      max_score: 12,
      has_zeros: response.has_zeros,
      verdict: response.verdict,
      dimensions,
    },
    snapshot,
  );
  const passed = policyDecision.passed;
  const verdict: JudgeResult["verdict"] = passed
    ? "PASS"
    : response.total_score >= 7
      ? "NEEDS_REVISION"
      : "REJECT";
  return {
    dimensions,
    total_score: response.total_score,
    max_score: 12,
    has_zeros: response.has_zeros,
    trajectory_pass: dimensions.trajectory_task_alignment.score > 0,
    dbdiff_pass: true,
    package_pass: passed,
    verdict,
    policy: {
      strictFullMarks,
      ignoreTaskComplexityForFullMarks: policyDecision.ignoreTaskComplexityForFullMarks,
      threshold: policyDecision.threshold,
      maxScore: policyDecision.maxScore,
      rule: policyDecision.rule,
    },
    summary: response.summary,
    revision_notes: response.revision_notes,
    dimensionLabels: { ...DIMENSION_LABELS },
    raw: response,
    evaluated: true,
    evaluatedAt: new Date().toISOString(),
    overallScore: response.total_score,
    maxScore: 12,
    passed,
    feedback: feedbackSections.filter(Boolean).join("\n\n"),
    rubricScores,
    revisionNotes: response.revision_notes,
    promptUsed,
    rawResponse,
  };
}

function resolveJudgeConfig(config: Config, snapshot: OptionsSnapshot): AiClientConfig {
  return {
    apiKey: config.judgeApiKey || config.defaultApiKey,
    baseURL: snapshot.judge?.baseURL || config.judgeBaseURL || config.defaultBaseURL,
    apiMode: "chat",
    model: snapshot.judge?.model || config.judgeModel || config.defaultModel,
    temperature: 0,
    enableThinking: false,
    clearThinking: false,
  };
}

export async function runJudge(input: RunJudgeInput): Promise<JudgeResult> {
  const { aiClient, config, snapshot, taskPackage, evidence } = input;
  const rubrics = Array.isArray(taskPackage.rubrics) ? taskPackage.rubrics : [];
  const groups = Array.isArray(taskPackage.subtasks) ? taskPackage.subtasks : [];
  const goldenTrajectory = resolveGoldenTrajectory(taskPackage, evidence, groups);
  let template = snapshot.judge?.promptTemplate || DEFAULT_JUDGE_PROMPT_TEMPLATE;
  if (!template.includes("${passPolicy}")) {
    template += "\n\n## Active Pass Policy\n${passPolicy}";
  }
  const unknownVariables = getUnknownTemplateVariables(template);
  if (unknownVariables.length > 0) {
    throw new Error("Prompt contains unsupported variables: " + unknownVariables.join(", "));
  }
  const prompt = buildJudgePrompt(template, {
    taskStatement: String(taskPackage.instruction || "").trim(),
    rubricJsonStr: buildRubricJson(rubrics),
    goldenTrajectory,
    passPolicy: buildPassPolicyText(snapshot),
  });

  const response = await aiClient.requestText(
    resolveJudgeConfig(config, snapshot),
    [
      { role: "system", content: JSON_ONLY_SYSTEM },
      { role: "user", content: prompt },
    ],
    {
      stream: false,
      jsonMode: false,
      jsonModeFallback: false,
      retries: 0,
      pluginCompat: true,
      disableBetaParameterFallback: true,
    },
  );
  const parsed = parseJudgeResponse(response.text);
  return mapJudgeResponseToResult(parsed, prompt, response, snapshot);
}

export { buildGoldenTrajectory };