"""Compare a candidate retrieval run with the baseline using ranx.

Exits with code 1 (the gate fails) if any metric drops by more than the tolerance. Exits with
code 2 for a setup error, such as a missing file or a run in the wrong shape.

Metrics
-------
    mrr@k, recall@k and ndcg@k (default k = 10), each compared on its own.

Tolerance
---------
    Default 0.02 per metric. This is a TEAM DEFAULT: measure the run-to-run variance in your
    sandbox before you lock it. See PB-RAG-02-04, section 4.3 (P3) and section 5.

Usage
-----
    python eval/eval_retrieval.py \\
        --candidate runs/candidate.json \\
        --baseline  baselines/baseline.json \\
        --qrels     baselines/qrels.json \\
        [--candidate-meta runs/candidate.meta.json] \\
        [--baseline-meta  baselines/baseline.meta.json] \\
        [--max-p95-increase 0.25] [--tolerance 0.02] [--k 10]

File formats
------------
    run (candidate and baseline):  {"<query_id>": {"<doc_id>": <score>, ...}, ...}
    qrels:                         {"<query_id>": {"<doc_id>": <relevance grade>, ...}, ...}
    Only questions that have qrels are scored. The run may hold more (for example unanswerable ones).
    meta (optional, latency):      {"queries": n, "search_latency_ms": {"p50": x, "p95": y}}
    The baseline run and its meta file come from export_run.py on the last accepted code.
    See PB-RAG-02-01 for how to build and sign off the golden set and the qrels.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Any, NoReturn

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger("eval_retrieval")

DEFAULT_TOLERANCE = 0.02  # team default: measure in your sandbox before locking it
DEFAULT_K = 10  # team default


def fail(message: str) -> NoReturn:
    log.error("SETUP ERROR: %s", message)
    sys.exit(2)


def load_json(path: Path, label: str) -> Any:
    if not path.exists():
        fail(f"{label} file not found at '{path}'. Complete the prerequisites in PB-RAG-02-04.")
    return json.loads(path.read_text(encoding="utf-8-sig"))


def validate_run(label: str, run: Any, qrels: dict[str, Any]) -> None:
    """Fail early, with a clear message, if a run is not {query: {doc: number}} or lacks queries."""
    if not isinstance(run, dict) or not run:
        fail(f"{label} must be a non-empty JSON object of query id -> document id -> score.")
    for query_id, documents in run.items():
        if not isinstance(documents, dict):
            fail(f"{label}: query '{query_id}' must map document ids to scores.")
        for doc_id, score in documents.items():
            if isinstance(score, bool) or not isinstance(score, (int, float)):
                fail(f"{label}: score for '{query_id}' / '{doc_id}' is not a number.")
    missing = [query_id for query_id in qrels if query_id not in run]
    if missing:
        fail(f"{label} has no results for {len(missing)} qrels question(s), for example {missing[:5]}.")


def score_run(qrels: dict[str, Any], run: dict[str, Any], metrics: list[str]) -> dict[str, float]:
    try:
        from ranx import Qrels, Run, evaluate
    except ImportError:
        fail("ranx is not installed. Run: pip install -r eval/requirements.txt")
    # ranx needs the run and the qrels to hold the same questions. Questions without qrels, such as
    # unanswerable ones, are not part of the retrieval metrics.
    result = evaluate(Qrels(qrels), Run({query_id: run[query_id] for query_id in qrels}), metrics)
    if isinstance(result, dict):
        return {name.lower(): float(value) for name, value in result.items()}
    return {metrics[0]: float(result)}


def p95(meta: dict[str, Any] | None) -> float | None:
    try:
        return float(meta["search_latency_ms"]["p95"]) if meta else None
    except (KeyError, TypeError, ValueError):
        return None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Compare a candidate run with the baseline.")
    parser.add_argument("--candidate", type=Path, required=True)
    parser.add_argument("--baseline", type=Path, required=True)
    parser.add_argument("--qrels", type=Path, default=Path("baselines/qrels.json"))
    parser.add_argument("--candidate-meta", type=Path)
    parser.add_argument("--baseline-meta", type=Path)
    parser.add_argument(
        "--max-p95-increase",
        type=float,
        help="fail if p95 search latency rises by more than this fraction (for example 0.25); "
        "without it, latency is reported only",
    )
    parser.add_argument("--tolerance", type=float, default=DEFAULT_TOLERANCE)
    parser.add_argument("--k", type=int, default=DEFAULT_K)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    metrics = [f"mrr@{args.k}", f"recall@{args.k}", f"ndcg@{args.k}"]

    qrels = load_json(args.qrels, "Qrels")
    baseline = load_json(args.baseline, "Baseline run")
    candidate = load_json(args.candidate, "Candidate run")
    validate_run("Qrels", qrels, {})
    validate_run("Baseline run", baseline, qrels)
    validate_run("Candidate run", candidate, qrels)

    baseline_scores = score_run(qrels, baseline, metrics)
    candidate_scores = score_run(qrels, candidate, metrics)

    regressions: list[str] = []
    log.info("%-14s %10s %10s %10s  %s", "Metric", "Baseline", "Candidate", "Delta", "Status")
    for metric in metrics:
        delta = candidate_scores[metric] - baseline_scores[metric]
        failed = delta < -args.tolerance
        log.info(
            "%-14s %10.4f %10.4f %+10.4f  %s",
            metric, baseline_scores[metric], candidate_scores[metric], delta,
            "FAIL" if failed else "PASS",
        )
        if failed:
            regressions.append(f"{metric}: {baseline_scores[metric]:.4f} -> {candidate_scores[metric]:.4f}")

    base_p95 = p95(load_json(args.baseline_meta, "Baseline meta") if args.baseline_meta else None)
    cand_p95 = p95(load_json(args.candidate_meta, "Candidate meta") if args.candidate_meta else None)
    if base_p95 and cand_p95:
        change = (cand_p95 - base_p95) / base_p95
        log.info("Search latency p95: %.1f ms -> %.1f ms (%+.0f%%)", base_p95, cand_p95, change * 100)
        if args.max_p95_increase is not None and change > args.max_p95_increase:
            regressions.append(f"search latency p95 rose {change:.0%}")

    if regressions:
        log.error(
            "GATE FAILED (tolerance %.4f per metric):\n  %s\n"
            "Next: triage with PB-RAG-02-03. Do not lower the tolerance without RAG Lead review. "
            "For an intentional baseline change, open a separate PR (PB-RAG-02-04, section 4.5).",
            args.tolerance, "\n  ".join(regressions),
        )
        sys.exit(1)
    log.info("GATE PASSED: every metric is within tolerance (%.4f).", args.tolerance)


if __name__ == "__main__":
    main()