from __future__ import annotations

import argparse
import glob
import json
import os
import sys
from collections import Counter
from pathlib import Path

EVAL_DIR = Path(__file__).resolve().parent
if str(EVAL_DIR) not in sys.path:
    sys.path.insert(0, str(EVAL_DIR))

from common.io import read_jsonl, write_jsonl  # noqa: E402
from common.loader import load_bench  # noqa: E402
from common.passk import default_k_list, print_report, summarize  # noqa: E402
from common.paths import ensure_sys_path  # noqa: E402

ensure_sys_path()


# Default judge per benchmark: all three use the LLM judge, matching the reference protocol.
BENCH_DEFAULT_JUDGE = {"aime": "llm", "hle": "llm", "gpqa": "llm"}


def load_raw(out_dir: str) -> list[dict]:
    """Merge every rank's raw output; for a given (task_id, sample_idx) the last write wins."""
    merged: dict[tuple, dict] = {}
    order: list[tuple] = []
    paths = sorted(glob.glob(os.path.join(out_dir, "raw_rank*.jsonl")))
    if not paths:
        raise FileNotFoundError(f"no raw_rank*.jsonl under {out_dir} (did generation finish?)")
    for path in paths:
        for row in read_jsonl(path):
            key = (row["task_id"], row["sample_idx"])
            if key not in merged:
                order.append(key)
            merged[key] = row
    return [merged[key] for key in order]


def _judge(results: list[dict], by_id: dict, bench_name: str, judge: str, concurrency: int) -> str | None:
    """Score `results` in place. Returns the judge model name for the LLM judge, else None."""
    if judge == "rule" and bench_name == "hle":
        raise SystemExit(
            "HLE does not support rule judging: the official protocol requires an LLM judge.\n"
            "Most HLE answers are free text, so normalized exact match marks equivalent-but-"
            "differently-written answers wrong and systematically understates the score.\n"
            "Set EVAL_JUDGE_BASE_URL / EVAL_JUDGE_API_KEY / EVAL_JUDGE_MODEL and use --judge llm."
        )
    if judge != "llm":
        for row in results:
            row["rule_correct"] = row["correct"]
        return None

    from common.judge import judge_answers, resolve_model, truncate_output

    judge_kind = "hle" if bench_name == "hle" else "equivalence"

    def judge_input(text: str) -> str:
        # HLE official protocol: drop the thinking section, send the rest in full, no truncation.
        if judge_kind == "hle":
            return text.split("</think>")[-1].strip()
        return truncate_output(text)

    items = [
        (
            by_id[row["task_id"]].get("judge_question") or by_id[row["task_id"]]["question"],
            by_id[row["task_id"]]["answer"],
            judge_input(row["raw_text"]),
        )
        for row in results
    ]
    for row, verdict in zip(results, judge_answers(items, kind=judge_kind, concurrency=concurrency), strict=True):
        row["rule_correct"] = row["correct"]
        row["correct"] = bool(verdict["judge_correct"])
        row["judge_raw"] = verdict.get("judge_raw", "")
        row["judge_parsed"] = bool(verdict.get("judge_parsed", True))
        if "judge_verdict_unparsed" in verdict:
            row["judge_verdict_unparsed"] = bool(verdict["judge_verdict_unparsed"])
        if verdict.get("judge_error"):
            row["judge_error"] = verdict["judge_error"]
    return resolve_model()


def score_run(
    out_dir: str,
    *,
    judge: str = "auto",
    judge_concurrency: int = 8,
    data_file: str | None = None,
    verbose: bool = True,
) -> int:
    meta_path = os.path.join(out_dir, "meta.json")
    meta: dict = {}
    if os.path.isfile(meta_path):
        with open(meta_path) as f:
            meta = json.load(f)

    bench_name = meta.get("bench", "aime")
    bench = load_bench(bench_name)
    if judge == "auto":
        judge = BENCH_DEFAULT_JUDGE.get(bench_name, "rule")
    data_path = data_file or meta.get("data_file")
    if not data_path:
        raise SystemExit(f"{meta_path} has no data_file; pass --data-file explicitly")
    problems = bench.load(data_path)
    records = bench.build_records(problems, limit=meta.get("limit"))
    by_id = {r["task_id"]: r for r in records}

    rollouts = load_raw(out_dir)
    unknown = [r["task_id"] for r in rollouts if r["task_id"] not in by_id]
    assert not unknown, f"raw output has task_ids that are not in the dataset: {sorted(set(unknown))[:5]}"

    results: list[dict] = []
    for rollout in rollouts:
        record = by_id[rollout["task_id"]]
        row = bench.score_rollout(rollout, record)
        results.append(row)

    judge_model = _judge(results, by_id, bench_name, judge, judge_concurrency)

    if not results:
        raise SystemExit(f"no rollout at all in {out_dir}/raw_rank*.jsonl; refusing to emit an empty score")

    # Coverage: missing rank files, problems or samples must be surfaced, never silently ignored.
    seen = {row["task_id"] for row in results}
    n_sample = int(meta.get("n_sample", max((row["sample_idx"] for row in results), default=0) + 1))
    missing = sorted(set(by_id) - seen)
    expected_rollouts = len(records) * n_sample
    coverage_complete = not missing and len(results) == expected_rollouts
    if not coverage_complete:
        print(
            f"[warn] incomplete coverage: {len(seen)}/{len(records)} tasks, {len(results)}/{expected_rollouts} "
            "rollouts; this is not a complete result",
            flush=True,
        )

    # Truncated rollouts and fallback extraction paths both depress the score, so surface them.
    finish_reasons = dict(Counter(str(row["finish_reason"]) for row in results if row.get("finish_reason")))
    extract_sources = dict(Counter(str(row["extract_source"]) for row in results if row.get("extract_source")))
    truncated = sum(finish_reasons.get(reason, 0) for reason in ("length", "prompt_length_exceeded"))
    invalid = sum(finish_reasons.get(reason, 0) for reason in ("ignored", "aborted", "error", "preempted"))

    summary = summarize(
        results,
        n_sample=n_sample,
        k_list=meta.get("k_list") or default_k_list(n_sample),
        extra={
            "bench": meta.get("bench", "aime"),
            "model_name": meta.get("model_name"),
            "judge": judge,
            "judge_model": judge_model,
            "temperature": meta.get("temperature"),
            "top_p": meta.get("top_p"),
            "reasoning_effort": meta.get("reasoning_effort"),
            "max_len": meta.get("max_len"),
            "extract_failed": sum(1 for row in results if row.get("extract_failed", row.get("pred") is None)) / len(results),
            "expected_tasks": len(records),
            "expected_rollouts": expected_rollouts,
            "coverage_complete": coverage_complete,
            "missing_tasks": len(missing),
            "ragged_samples": len({sum(1 for row in results if row["task_id"] == task_id) for task_id in seen}) > 1,
            "finish_reasons": finish_reasons,
            "extract_sources": extract_sources,
            "truncated_rollouts": truncated,
            "invalid_rollouts": invalid,
        },
    )
    groups: dict[str, list[dict]] = {}
    for row in results:
        if key := str(by_id[row["task_id"]].get("modality") or ""):
            groups.setdefault(key, []).append(row)
    if len(groups) > 1:
        summary["by_modality"] = {
            modality: {
                "num_tasks": len({row["task_id"] for row in rows}),
                "num_rollouts": len(rows),
                "Pass@1": sum(1 for row in rows if row["correct"]) / len(rows),
            }
            for modality, rows in sorted(groups.items())
        }

    if judge == "llm":
        agree = sum(1 for row in results if bool(row["rule_correct"]) == bool(row["correct"]))
        summary["rule_vs_llm_agreement"] = agree / len(results)
        unparsed = sum(1 for row in results if not row.get("judge_parsed", True))
        summary["judge_unparsed"] = unparsed / len(results)
        verdict_unparsed = sum(1 for row in results if row.get("judge_verdict_unparsed"))
        summary["judge_verdict_unparsed"] = verdict_unparsed / len(results)
        if verdict_unparsed:
            print(
                f"[warn] {verdict_unparsed}/{len(results)} judge replies had no extractable verdict; "
                "all three fallbacks missed, they count as wrong",
                flush=True,
            )
        if unparsed:
            print(
                f"[warn] {unparsed}/{len(results)} judge calls returned no verdict and count as wrong; "
                "this depresses the score -- raise the judge max_tokens and re-run",
                flush=True,
            )

    os.makedirs(out_dir, exist_ok=True)
    write_jsonl(os.path.join(out_dir, "scored.jsonl"), results)
    with open(os.path.join(out_dir, "summary.json"), "w") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)

    if verbose:
        print_report(summary, title=f"{summary['bench']}  {summary.get('model_name')}  judge={judge}")
        print(
            f"coverage: expected {summary['expected_rollouts']} rollouts / "
            f"{summary['expected_tasks']} tasks, got {summary['num_rollouts']} / {summary['num_tasks']}, "
            f"complete={summary['coverage_complete']}"
        )
        print(f"finish_reason={summary['finish_reasons']}")
        for key, stat in (summary.get("by_modality") or {}).items():
            print(f"  Pass@1[{key}] = {stat['Pass@1'] * 100:6.2f}  (tasks={stat['num_tasks']})")
        if summary["extract_sources"]:
            print(f"extract_source={summary['extract_sources']}")
        if summary["truncated_rollouts"] or summary["invalid_rollouts"]:
            print(
                f"[warn] {summary['truncated_rollouts']} truncated / "
                f"{summary['invalid_rollouts']} invalid rollouts, all counted as wrong; "
                "this depresses the score"
            )
        print(f"extract_failed={summary['extract_failed'] * 100:.2f}%")
        if "rule_vs_llm_agreement" in summary:
            print(f"rule vs llm agreement = {summary['rule_vs_llm_agreement'] * 100:.2f}%")
        print(f"wrote {os.path.join(out_dir, 'scored.jsonl')} / summary.json")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="score a benchmark run and summarize Pass@k")
    parser.add_argument("--out-dir", required=True)
    parser.add_argument(
        "--judge", choices=["auto", "rule", "llm"], default="auto", help="auto: per-benchmark default (all three use the LLM judge)"
    )
    parser.add_argument("--judge-concurrency", type=int, default=8)
    parser.add_argument("--data-file", default=None, help="override meta.json's data_file")
    args = parser.parse_args()
    return score_run(
        args.out_dir,
        judge=args.judge,
        judge_concurrency=args.judge_concurrency,
        data_file=args.data_file,
    )


if __name__ == "__main__":
    raise SystemExit(main())