"""Re-Judge MCP — Re-evaluate existing generations with alternative pairwise prompts.

Usage:
    python -m evaluation.rejudge --run-dir evaluation/results/run_XXXXX --variant discovery_aware
    python -m evaluation.rejudge --run-dir evaluation/results/run_XXXXX --variant balanced
    python -m evaluation.rejudge --list-variants
    python -m evaluation.rejudge --run-dir evaluation/results/run_XXXXX --variant custom --custom-prompt path/to/prompt.txt

Built-in prompt variants:
    original           — Current pairwise prompt (baseline for comparison)
    discovery_aware    — Adds "Discovery" as 6th criterion
    discovery_weighted — Discovery as top-priority criterion
    error_first        — Error absence as highest weight
    balanced           — All 6 criteria equally weighted, no hierarchy
    integration_only   — ONLY integration criterion (diagnostic)
    custom             — Load from --custom-prompt file
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import random
import re
import sys
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from .config import EXPERIMENT_CONFIG
from .corpus.loader import TestText, load_corpus
from .judges.llm_client import JudgeLLM
from .judges.prompts import PAIRWISE_USER_PROMPT
from .sanitize import sanitize_for_judge
from .stats.bradley_terry import fit_bradley_terry, strengths_to_elo

logger = logging.getLogger(__name__)

CONDITIONS = EXPERIMENT_CONFIG["conditions"]
SEED = EXPERIMENT_CONFIG["seed"]


# ══════════════════════════════════════════════════════════════════════
# Prompt Variants
# ══════════════════════════════════════════════════════════════════════

PROMPT_VARIANTS: Dict[str, str] = {}

# ── Original (baseline) ──────────────────────────────────────────────
PROMPT_VARIANTS["original"] = """\
You are a blind evaluator comparing two analyses of the same text.
You do NOT know which analysis was produced by which method.

EVALUATION CRITERIA (in order of importance):
1. Structural depth: Does the analysis reveal hidden structure in the text?
2. Multi-perspective coverage: How many distinct analytical angles are applied?
3. Integration: Do the perspectives connect to each other meaningfully?
4. Insight quality: Are the insights non-obvious and well-grounded?
5. Error absence: Is the analysis free from logical errors and unsupported claims?

INSTRUCTIONS:
- Read both analyses carefully
- Compare on each criterion above
- You MUST pick a winner. There is NO tie option.
- Even if both analyses are close in quality, identify the one that is \
marginally better. There is always a difference.
- If criteria votes are split, use this tie-breaking hierarchy: \
Insight quality > Structural depth > Integration > Multi-perspective > Error absence

OUTPUT FORMAT:
Winner: [A or B]
Reasoning: [2-3 sentences explaining the decision]
Per-criterion:
  Structural depth: [A or B]
  Multi-perspective: [A or B]
  Integration: [A or B]
  Insight quality: [A or B]
  Error absence: [A or B]
"""

# ── Discovery-Aware ──────────────────────────────────────────────────
PROMPT_VARIANTS["discovery_aware"] = """\
You are a blind evaluator comparing two analyses of the same text.
You do NOT know which analysis was produced by which method.

EVALUATION CRITERIA (in order of importance):
1. Structural depth: Does the analysis reveal hidden structure in the text?
2. Multi-perspective coverage: How many distinct analytical angles are applied?
3. Integration: Do the perspectives connect to each other meaningfully?
4. Insight quality: Are the insights non-obvious and well-grounded?
5. Error absence: Is the analysis free from logical errors and unsupported claims?
6. Discovery: Does the analysis surface genuinely non-obvious connections or \
insights that would surprise an educated reader? Look for unexpected cross-domain \
links, hidden patterns, or novel interpretations grounded in the original text.

INSTRUCTIONS:
- Read both analyses carefully
- Compare on each criterion above
- You MUST pick a winner. There is NO tie option.
- Even if both analyses are close in quality, identify the one that is \
marginally better. There is always a difference.
- If criteria votes are split, use this tie-breaking hierarchy: \
Discovery > Insight quality > Structural depth > Integration > Multi-perspective > Error absence

OUTPUT FORMAT:
Winner: [A or B]
Reasoning: [2-3 sentences explaining the decision]
Per-criterion:
  Structural depth: [A or B]
  Multi-perspective: [A or B]
  Integration: [A or B]
  Insight quality: [A or B]
  Error absence: [A or B]
  Discovery: [A or B]
"""

# ── Discovery-Weighted (discovery as TOP criterion) ──────────────────
PROMPT_VARIANTS["discovery_weighted"] = """\
You are a blind evaluator comparing two analyses of the same text.
You do NOT know which analysis was produced by which method.

EVALUATION CRITERIA (in order of importance):
1. Discovery: Does the analysis surface genuinely non-obvious connections or \
insights that would surprise an educated reader? Look for unexpected cross-domain \
links, hidden patterns, novel interpretations grounded in the original text. \
This is the MOST important criterion.
2. Insight quality: Are the insights non-obvious and well-grounded?
3. Structural depth: Does the analysis reveal hidden structure in the text?
4. Integration: Do the perspectives connect to each other meaningfully?
5. Error absence: Is the analysis free from logical errors and unsupported claims?
6. Multi-perspective coverage: How many distinct analytical angles are applied?

INSTRUCTIONS:
- Read both analyses carefully
- Compare on each criterion above
- You MUST pick a winner. There is NO tie option.
- Even if both analyses are close in quality, identify the one that is \
marginally better. There is always a difference.
- If criteria votes are split, use this tie-breaking hierarchy: \
Discovery > Insight quality > Structural depth > Integration > Error absence > Multi-perspective

OUTPUT FORMAT:
Winner: [A or B]
Reasoning: [2-3 sentences explaining the decision]
Per-criterion:
  Discovery: [A or B]
  Insight quality: [A or B]
  Structural depth: [A or B]
  Integration: [A or B]
  Error absence: [A or B]
  Multi-perspective: [A or B]
"""

# ── Error-First ──────────────────────────────────────────────────────
PROMPT_VARIANTS["error_first"] = """\
You are a blind evaluator comparing two analyses of the same text.
You do NOT know which analysis was produced by which method.

EVALUATION CRITERIA (in order of importance):
1. Error absence: Is the analysis free from logical errors, unsupported claims, \
circular reasoning, hallucinated content, and contradictions? This is the MOST \
important criterion — accuracy trumps all else.
2. Insight quality: Are the insights non-obvious and well-grounded?
3. Structural depth: Does the analysis reveal hidden structure in the text?
4. Integration: Do the perspectives connect to each other meaningfully?
5. Discovery: Does the analysis surface genuinely non-obvious connections?
6. Multi-perspective coverage: How many distinct analytical angles are applied?

INSTRUCTIONS:
- Read both analyses carefully
- Compare on each criterion above
- You MUST pick a winner. There is NO tie option.
- Even if both analyses are close in quality, identify the one that is \
marginally better. There is always a difference.
- If criteria votes are split, use this tie-breaking hierarchy: \
Error absence > Insight quality > Structural depth > Integration > Discovery > Multi-perspective

OUTPUT FORMAT:
Winner: [A or B]
Reasoning: [2-3 sentences explaining the decision]
Per-criterion:
  Error absence: [A or B]
  Insight quality: [A or B]
  Structural depth: [A or B]
  Integration: [A or B]
  Discovery: [A or B]
  Multi-perspective: [A or B]
"""

# ── Balanced (no hierarchy, all equal) ───────────────────────────────
PROMPT_VARIANTS["balanced"] = """\
You are a blind evaluator comparing two analyses of the same text.
You do NOT know which analysis was produced by which method.

EVALUATION CRITERIA (all equally important — no criterion outranks another):
- Structural depth: Does the analysis reveal hidden structure in the text?
- Multi-perspective coverage: How many distinct analytical angles are applied?
- Integration: Do the perspectives connect to each other meaningfully?
- Insight quality: Are the insights non-obvious and well-grounded?
- Error absence: Is the analysis free from logical errors and unsupported claims?
- Discovery: Does the analysis surface genuinely non-obvious connections or \
insights that would surprise an educated reader?

INSTRUCTIONS:
- Read both analyses carefully
- Compare on EACH criterion above — they are all equally weighted
- Count how many criteria each analysis wins
- The analysis that wins MORE criteria is the overall winner
- You MUST pick a winner. There is NO tie option. If tied on criteria count, \
re-read both analyses holistically and pick the one with the strongest \
single-criterion advantage.

OUTPUT FORMAT:
Winner: [A or B]
Reasoning: [2-3 sentences explaining the decision]
Per-criterion:
  Structural depth: [A or B]
  Multi-perspective: [A or B]
  Integration: [A or B]
  Insight quality: [A or B]
  Error absence: [A or B]
  Discovery: [A or B]
"""

# ── Integration-Only (diagnostic) ────────────────────────────────────
PROMPT_VARIANTS["integration_only"] = """\
You are a blind evaluator comparing two analyses of the same text.
You do NOT know which analysis was produced by which method.

EVALUATION CRITERION (only one):
Integration: Do the perspectives connect to each other meaningfully? \
Does understanding from one lens inform or constrain another? \
Is the analysis greater than the sum of its parts?

INSTRUCTIONS:
- Read both analyses carefully
- Focus ONLY on how well different analytical perspectives are integrated
- Ignore all other quality dimensions (depth, discovery, errors, etc.)
- You MUST pick a winner. There is NO tie option.

OUTPUT FORMAT:
Winner: [A or B]
Reasoning: [2-3 sentences explaining the decision]
Per-criterion:
  Integration: [A or B]
"""


# ══════════════════════════════════════════════════════════════════════
# Pairwise comparison (with custom prompt)
# ══════════════════════════════════════════════════════════════════════

@dataclass
class PairwiseResult:
    text_id: str
    condition_a: str
    condition_b: str
    winner: str
    true_winner: str
    reasoning: str = ""
    per_criterion: Dict[str, str] = field(default_factory=dict)


def _parse_pairwise(response: str, variant: str) -> Tuple[str, str, Dict[str, str]]:
    """Parse winner, reasoning, per-criterion from judge response."""
    m = re.search(r"Winner:\s*(A|B|TIE)", response, re.IGNORECASE)
    winner = m.group(1).upper() if m else "TIE"
    if winner == "TIE":
        logger.warning("Judge returned TIE despite forced-choice; breaking randomly")
        winner = random.choice(["A", "B"])

    m = re.search(r"Reasoning:\s*(.+?)(?:\n|Per-criterion)", response, re.DOTALL)
    reasoning = m.group(1).strip() if m else ""

    criteria = {}
    possible_criteria = [
        "Structural depth", "Multi-perspective", "Integration",
        "Insight quality", "Error absence", "Discovery",
    ]
    for crit in possible_criteria:
        m = re.search(rf"{crit}:\s*(A|B|TIE)", response, re.IGNORECASE)
        if m:
            criteria[crit] = m.group(1).upper()
    return winner, reasoning, criteria


async def compare_pair(
    judge: JudgeLLM,
    system_prompt: str,
    text_id: str,
    input_text: str,
    output_a: str,
    condition_a: str,
    output_b: str,
    condition_b: str,
    variant: str,
) -> PairwiseResult:
    """One pairwise comparison with randomised presentation order."""
    if random.random() < 0.5:
        shown_a, shown_b = output_a, output_b
        map_a, map_b = condition_a, condition_b
    else:
        shown_a, shown_b = output_b, output_a
        map_a, map_b = condition_b, condition_a

    response = await judge.judge(
        system_prompt,
        PAIRWISE_USER_PROMPT.format(
            input_text=input_text[:2000],
            analysis_a=shown_a[:4000],
            analysis_b=shown_b[:4000],
        ),
    )

    winner_label, reasoning, per_criterion = _parse_pairwise(response, variant)

    if winner_label == "A":
        true_winner = map_a
    elif winner_label == "B":
        true_winner = map_b
    else:
        true_winner = "TIE"

    return PairwiseResult(
        text_id=text_id,
        condition_a=map_a,
        condition_b=map_b,
        winner=winner_label,
        true_winner=true_winner,
        reasoning=reasoning,
        per_criterion=per_criterion,
    )


async def run_all_pairwise_variant(
    judge: JudgeLLM,
    system_prompt: str,
    text_id: str,
    input_text: str,
    outputs: Dict[str, str],
    variant: str,
) -> List[PairwiseResult]:
    """Run all C(n,2) pairwise comparisons for one text."""
    conditions = list(outputs.keys())
    results: List[PairwiseResult] = []
    for i in range(len(conditions)):
        for j in range(i + 1, len(conditions)):
            result = await compare_pair(
                judge, system_prompt, text_id, input_text,
                outputs[conditions[i]], conditions[i],
                outputs[conditions[j]], conditions[j],
                variant,
            )
            results.append(result)
    return results


# ══════════════════════════════════════════════════════════════════════
# Report generation
# ══════════════════════════════════════════════════════════════════════

def _load_old_elo(run_dir: Path) -> Dict[str, float]:
    elo_path = run_dir / "elo_ratings.json"
    if elo_path.exists():
        return json.loads(elo_path.read_text(encoding="utf-8"))
    return {}


def _pairwise_to_records(pairwise: List[PairwiseResult]) -> List[Dict]:
    records: List[Dict] = []
    for r in pairwise:
        if r.true_winner == "TIE":
            records.append({"winner": "TIE", "cond_a": r.condition_a, "cond_b": r.condition_b})
        else:
            loser = r.condition_b if r.true_winner == r.condition_a else r.condition_a
            records.append({"winner": r.true_winner, "loser": loser,
                            "cond_a": r.condition_a, "cond_b": r.condition_b})
    return records


def _win_counts(pairwise: List[PairwiseResult]) -> Dict[str, int]:
    counts: Dict[str, int] = {c: 0 for c in CONDITIONS}
    for r in pairwise:
        if r.true_winner in counts:
            counts[r.true_winner] += 1
    return counts


def _criterion_wins(pairwise: List[PairwiseResult], c1: str, c2: str) -> Dict[str, Dict[str, int]]:
    """Per-criterion wins between two conditions."""
    crits: Dict[str, Dict[str, int]] = {}
    for r in pairwise:
        if set([r.condition_a, r.condition_b]) != set([c1, c2]):
            continue
        for crit, vote in r.per_criterion.items():
            if crit not in crits:
                crits[crit] = {c1: 0, c2: 0}
            if vote == "A":
                crits[crit][r.condition_a] += 1
            elif vote == "B":
                crits[crit][r.condition_b] += 1
    return crits


def _domain_difficulty_breakdown(
    pairwise: List[PairwiseResult],
    corpus: List[TestText],
    c1: str, c2: str,
) -> Tuple[Dict[str, Dict[str, int]], Dict[str, Dict[str, int]]]:
    """Win counts by domain and difficulty for a pair of conditions."""
    text_meta = {t.id: t for t in corpus}
    by_domain: Dict[str, Dict[str, int]] = {}
    by_diff: Dict[str, Dict[str, int]] = {}

    for r in pairwise:
        if set([r.condition_a, r.condition_b]) != set([c1, c2]):
            continue
        meta = text_meta.get(r.text_id)
        if not meta:
            continue

        for key, bucket in [(meta.domain, by_domain), (meta.difficulty, by_diff)]:
            if key not in bucket:
                bucket[key] = {c1: 0, c2: 0}
            if r.true_winner in bucket[key]:
                bucket[key][r.true_winner] += 1

    return by_domain, by_diff


def generate_rejudge_report(
    *,
    run_dir: Path,
    out_dir: Path,
    variant: str,
    corpus: List[TestText],
    pairwise: List[PairwiseResult],
    elo_new: Dict[str, float],
    elo_old: Dict[str, float],
) -> Path:
    """Generate comparison report: old prompt vs new prompt."""
    lines: List[str] = []
    a = lines.append

    a(f"# Re-Judge Report: `{variant}`")
    a("")
    a(f"**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M')}")
    a(f"**Source run:** `{run_dir.name}`")
    a(f"**Prompt variant:** `{variant}`")
    a(f"**Corpus:** {len(corpus)} texts")
    a("")

    # ── Elo comparison ────────────────────────────────────────────
    a("## Elo Comparison")
    a("")
    a("| Condition | Original Elo | New Elo | Delta |")
    a("|-----------|:----------:|:------:|:-----:|")
    for cond in CONDITIONS:
        old = elo_old.get(cond, 0)
        new = elo_new.get(cond, 0)
        delta = new - old
        sign = "+" if delta >= 0 else ""
        a(f"| {cond} | {old:.0f} | {new:.0f} | {sign}{delta:.0f} |")
    a("")

    # ── Win counts ────────────────────────────────────────────────
    wins = _win_counts(pairwise)
    a("## Pairwise Win Counts")
    a("")
    a("| Condition | Wins |")
    a("|-----------|:----:|")
    for cond in CONDITIONS:
        a(f"| {cond} | {wins.get(cond, 0)} |")
    a("")

    # ── LLM vs VOCAB per-criterion ────────────────────────────────
    a("## LLM_SCAFFOLD vs VOCAB_ONLY — Per-Criterion")
    a("")
    crit_wins = _criterion_wins(pairwise, "LLM_SCAFFOLD", "VOCAB_ONLY")
    if crit_wins:
        a("| Criterion | LLM_SCAFFOLD | VOCAB_ONLY |")
        a("|-----------|:----------:|:----------:|")
        for crit, votes in sorted(crit_wins.items()):
            a(f"| {crit} | {votes.get('LLM_SCAFFOLD', 0)} | {votes.get('VOCAB_ONLY', 0)} |")
        a("")

    # ── Domain breakdown ──────────────────────────────────────────
    a("## LLM vs VOCAB — By Domain")
    a("")
    by_domain, by_diff = _domain_difficulty_breakdown(
        pairwise, corpus, "LLM_SCAFFOLD", "VOCAB_ONLY"
    )
    if by_domain:
        a("| Domain | LLM | VOCAB |")
        a("|--------|:---:|:-----:|")
        for dom in sorted(by_domain):
            a(f"| {dom} | {by_domain[dom].get('LLM_SCAFFOLD', 0)} | {by_domain[dom].get('VOCAB_ONLY', 0)} |")
        a("")

    # ── Difficulty breakdown ──────────────────────────────────────
    a("## LLM vs VOCAB — By Difficulty")
    a("")
    if by_diff:
        a("| Difficulty | LLM | VOCAB |")
        a("|------------|:---:|:-----:|")
        for diff in ["easy", "medium", "hard"]:
            if diff in by_diff:
                a(f"| {diff} | {by_diff[diff].get('LLM_SCAFFOLD', 0)} | {by_diff[diff].get('VOCAB_ONLY', 0)} |")
        a("")

    # ── NAKED vs others ───────────────────────────────────────────
    a("## NAKED vs Others")
    a("")
    for other in ["VOCAB_ONLY", "LLM_SCAFFOLD"]:
        wins_nk = sum(1 for r in pairwise
                      if set([r.condition_a, r.condition_b]) == set(["NAKED", other])
                      and r.true_winner == "NAKED")
        wins_ot = sum(1 for r in pairwise
                      if set([r.condition_a, r.condition_b]) == set(["NAKED", other])
                      and r.true_winner == other)
        a(f"- NAKED vs {other}: {wins_nk} — {wins_ot}")
    a("")

    # ── Prompt used ───────────────────────────────────────────────
    a("## Prompt Used")
    a("")
    a("```")
    a(PROMPT_VARIANTS.get(variant, "custom"))
    a("```")

    report_path = out_dir / "rejudge_report.md"
    report_path.write_text("\n".join(lines), encoding="utf-8")
    return report_path


# ══════════════════════════════════════════════════════════════════════
# Main pipeline
# ══════════════════════════════════════════════════════════════════════

async def rejudge(
    run_dir: Path,
    variant: str,
    custom_prompt: Optional[str] = None,
    subset_domain: Optional[str] = None,
    subset_difficulty: Optional[str] = None,
) -> Path:
    """Re-judge an existing run with a different pairwise prompt."""
    random.seed(SEED)
    np.random.seed(SEED)

    # Resolve prompt
    if variant == "custom":
        if not custom_prompt:
            raise ValueError("--custom-prompt required for 'custom' variant")
        system_prompt = Path(custom_prompt).read_text(encoding="utf-8")
        PROMPT_VARIANTS["custom"] = system_prompt
    else:
        if variant not in PROMPT_VARIANTS:
            raise ValueError(f"Unknown variant: {variant}. Use --list-variants")
        system_prompt = PROMPT_VARIANTS[variant]

    # Output directory
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    out_dir = run_dir / f"rejudge_{variant}_{timestamp}"
    out_dir.mkdir(parents=True, exist_ok=True)

    # Save config
    config = {
        "source_run": str(run_dir),
        "variant": variant,
        "timestamp": timestamp,
        "subset_domain": subset_domain,
        "subset_difficulty": subset_difficulty,
    }
    (out_dir / "config.json").write_text(
        json.dumps(config, indent=2), encoding="utf-8"
    )

    # Load corpus
    corpus = load_corpus()
    if subset_domain:
        corpus = [t for t in corpus if t.domain == subset_domain]
    if subset_difficulty:
        corpus = [t for t in corpus if t.difficulty == subset_difficulty]
    logger.info("Corpus: %d texts (domain=%s, difficulty=%s)",
                len(corpus), subset_domain or "all", subset_difficulty or "all")

    # Load existing raw outputs
    raw_dir = run_dir / "raw"
    if not raw_dir.exists():
        raise FileNotFoundError(f"No raw outputs at {raw_dir}")

    outputs: Dict[str, Dict[str, str]] = {}
    for text in corpus:
        outputs[text.id] = {}
        for cond_id in CONDITIONS:
            fp = raw_dir / cond_id / f"{text.id}.json"
            if not fp.exists():
                logger.warning("Missing raw output: %s", fp)
                continue
            data = json.loads(fp.read_text(encoding="utf-8"))
            outputs[text.id][cond_id] = sanitize_for_judge(
                data["raw_output"], cond_id
            )

    # Run pairwise with new prompt
    judge = JudgeLLM()
    all_pairwise: List[PairwiseResult] = []
    total = len(corpus)

    for idx, text in enumerate(corpus, 1):
        logger.info("[%d/%d] Pairwise: %s", idx, total, text.id)
        if text.id not in outputs or len(outputs[text.id]) < 2:
            logger.warning("Skipping %s — insufficient outputs", text.id)
            continue
        results = await run_all_pairwise_variant(
            judge, system_prompt, text.id, text.text,
            outputs[text.id], variant,
        )
        all_pairwise.extend(results)

    # Save pairwise results
    pw_path = out_dir / "all_comparisons.json"
    pw_path.write_text(json.dumps(
        [vars(r) for r in all_pairwise], indent=2, ensure_ascii=False,
    ), encoding="utf-8")

    # Bradley-Terry → Elo
    bt_records = _pairwise_to_records(all_pairwise)
    bt_strengths = fit_bradley_terry(bt_records, CONDITIONS)
    elo_new = strengths_to_elo(bt_strengths)

    elo_path = out_dir / "elo_ratings.json"
    elo_path.write_text(json.dumps(elo_new, indent=2), encoding="utf-8")

    # Load old Elo for comparison
    elo_old = _load_old_elo(run_dir)

    # Generate report
    report_path = generate_rejudge_report(
        run_dir=run_dir,
        out_dir=out_dir,
        variant=variant,
        corpus=corpus,
        pairwise=all_pairwise,
        elo_new=elo_new,
        elo_old=elo_old,
    )

    logger.info("Re-judge complete → %s", report_path)
    logger.info("New Elo: %s", {k: f"{v:.0f}" for k, v in elo_new.items()})
    if elo_old:
        for cond in CONDITIONS:
            delta = elo_new.get(cond, 0) - elo_old.get(cond, 0)
            logger.info("  %s: %+.0f", cond, delta)

    return report_path


# ══════════════════════════════════════════════════════════════════════
# CLI
# ══════════════════════════════════════════════════════════════════════

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Re-Judge: re-evaluate existing run with alternative pairwise prompts"
    )
    parser.add_argument("--run-dir", type=str,
                        help="Path to existing run directory")
    parser.add_argument("--variant", type=str, default="discovery_aware",
                        choices=list(PROMPT_VARIANTS.keys()) + ["custom"],
                        help="Prompt variant to use")
    parser.add_argument("--custom-prompt", type=str,
                        help="Path to custom prompt file (for --variant custom)")
    parser.add_argument("--domain", type=str, default=None,
                        help="Filter to specific domain (e.g. ETH, SCI)")
    parser.add_argument("--difficulty", type=str, default=None,
                        help="Filter to specific difficulty (easy, medium, hard)")
    parser.add_argument("--list-variants", action="store_true",
                        help="List available prompt variants and exit")
    args = parser.parse_args()

    if args.list_variants:
        print("Available prompt variants:\n")
        for name, prompt in PROMPT_VARIANTS.items():
            first_line = prompt.strip().split("\n")[0]
            print(f"  {name:25s} — {first_line[:70]}")
        return

    if not args.run_dir:
        parser.error("--run-dir is required")

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    run_dir = Path(args.run_dir)
    if not run_dir.exists():
        print(f"ERROR: Run directory not found: {run_dir}", file=sys.stderr)
        sys.exit(1)

    asyncio.run(rejudge(
        run_dir=run_dir,
        variant=args.variant,
        custom_prompt=args.custom_prompt,
        subset_domain=args.domain,
        subset_difficulty=args.difficulty,
    ))


if __name__ == "__main__":
    main()
