#!/usr/bin/env python3
"""Retrieval evaluation harness for QModel.

Runs every case in a golden set through the retrieval stack (no LLM generation)
and reports ranking metrics plus confidence-gate behaviour.

Golden set format (JSONL, one case per line):
    {"id": "...", "category": "...", "lang": "ar|en", "query": "...", "relevant": ["2:255", "bukhari_1", ...]}
`relevant` holds dataset `id`s; any one of them counts as a hit. An empty list
marks a negative (out-of-scope) case that the confidence gate should reject.
Optional `required` (a subset of `relevant`) names sources a complete answer must
include — e.g. 5:90 for "is alcohol forbidden?", where 2:219 alone is a hit but
misleading. These are scored separately as complete@k, so hit@k stays comparable.

Modes:
    pipeline  — the full pre-generation pipeline (_rag_preprocess): rewrite, hybrid
                search, direct lookups, text search, merging, dedup. This is exactly
                what the LLM receives as context.
    hybrid    — hybrid_search() alone (dense + BM25 + RRF [+ reranker]), useful for
                isolating fusion / embedding changes.

Usage:
    python -m eval.run_eval                         # pipeline mode, k = TOP_K_RETURN
    python -m eval.run_eval --mode hybrid --k 20
    python -m eval.run_eval --only fatwa,negative --show-failures
    python -m eval.run_eval --out eval/results/baseline.json
    python -m eval.run_eval --compare eval/results/baseline.json
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import re
import subprocess
import sys
import time
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from app.config import cfg  # noqa: E402

DEFAULT_GOLDEN = ROOT / "eval" / "golden.jsonl"
RESULTS_DIR = ROOT / "eval" / "results"
_GRADE_SUFFIX = re.compile(r"\s*\[[^\]]*\]\s*$")


# ═══════════════════════════════════════════════════════════════════════
# LOADING
# ═══════════════════════════════════════════════════════════════════════
def load_golden(path: Path, only: Optional[set]) -> List[dict]:
    cases = []
    with open(path, encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            line = line.strip()
            if not line or line.startswith("//"):
                continue
            case = json.loads(line)
            for key in ("id", "category", "query", "relevant"):
                if key not in case:
                    raise ValueError(f"{path}:{line_no}: missing '{key}'")
            extra = set(case.get("required", [])) - set(case["relevant"])
            if extra:
                raise ValueError(f"{path}:{line_no}: required ids not in relevant: {sorted(extra)}")
            if only and case["category"] not in only:
                continue
            cases.append(case)
    ids = [c["id"] for c in cases]
    dupes = {i for i in ids if ids.count(i) > 1}
    if dupes:
        raise ValueError(f"Duplicate case ids: {sorted(dupes)}")
    return cases


def init_state(use_llm: bool) -> None:
    """Load models/index into app.state the same way the server lifespan does."""
    import faiss
    from sentence_transformers import SentenceTransformer

    from app.search import build_type_indexes
    from app.state import _build_lookup_tables, infer_hadith_grade, state

    t0 = time.perf_counter()
    print(f"Loading embed model {cfg.EMBED_MODEL} …", flush=True)
    state.embed_model = SentenceTransformer(cfg.EMBED_MODEL)
    print(f"Loading FAISS index {cfg.FAISS_INDEX} …", flush=True)
    state.faiss_index = faiss.read_index(str(ROOT / cfg.FAISS_INDEX))
    print(f"Loading metadata {cfg.METADATA_FILE} …", flush=True)
    with open(ROOT / cfg.METADATA_FILE, encoding="utf-8") as f:
        state.dataset = [infer_hadith_grade(item) for item in json.load(f)]
    print("Building lookup tables + BM25 …", flush=True)
    _build_lookup_tables(state.dataset)
    state.type_indexes = build_type_indexes(state.faiss_index, state.dataset)

    state.reranker = None
    if cfg.USE_RERANKER:
        from sentence_transformers import CrossEncoder
        print(f"Loading reranker {cfg.RERANK_MODEL} …", flush=True)
        state.reranker = CrossEncoder(cfg.RERANK_MODEL, max_length=512)

    state.llm = None
    if use_llm:
        from app.llm import get_llm_provider
        print(f"Initializing LLM provider {cfg.LLM_BACKEND} …", flush=True)
        state.llm = get_llm_provider()
    elif not cfg.FAST_QUERY_REWRITE:
        print("WARNING: FAST_QUERY_REWRITE is off but --llm not given; "
              "queries will use the regex fallback rewrite.", flush=True)

    state.ready = True
    print(f"Ready in {time.perf_counter() - t0:.1f}s | dataset={len(state.dataset):,} "
          f"| faiss={state.faiss_index.ntotal:,} | reranker={'on' if state.reranker else 'off'}\n",
          flush=True)


def validate_labels(cases: List[dict], dataset: list) -> None:
    known = {item.get("id") for item in dataset}
    missing = [(c["id"], r) for c in cases for r in c["relevant"] if r not in known]
    if missing:
        lines = "\n".join(f"  {cid}: {rid}" for cid, rid in missing[:20])
        raise SystemExit(f"Golden set references {len(missing)} unknown dataset ids:\n{lines}")


# ═══════════════════════════════════════════════════════════════════════
# RETRIEVAL
# ═══════════════════════════════════════════════════════════════════════
async def retrieve(case: dict, mode: str, k: int) -> dict:
    from app.search import hybrid_search, rewrite_query
    from app.state import _rag_preprocess, state

    if mode == "pipeline":
        ctx = await _rag_preprocess(case["query"], k, None, None)
        return {
            "results": ctx["results"],
            "intent": ctx["intent"],
            "top_score": ctx["top_score"],
            "gated": not ctx["confident"],
            "confidence": ctx["confidence"],
            "confidence_source": ctx["confidence_source"],
        }

    rewrite = await rewrite_query(case["query"], state.llm)
    results = await hybrid_search(
        case["query"], rewrite, state.embed_model, state.faiss_index, state.dataset,
        top_n=k, bm25_index=state.bm25_index, reranker_model=state.reranker,
    )
    # The gate runs in _rag_preprocess, so hybrid mode doesn't evaluate it.
    return {
        "results": results,
        "intent": rewrite.get("intent", "general"),
        "top_score": results[0].get("_score", 0.0) if results else 0.0,
        "gated": None,
        "confidence": None,
        "confidence_source": None,
    }


def result_ids(item: dict, ref_to_id: Dict[str, str]) -> set:
    """All dataset ids a result stands for, including narrations merged into it by dedup."""
    ids = {item.get("id")}
    for alt in item.get("other_collections") or []:
        ref = _GRADE_SUFFIX.sub("", alt).strip()
        if ref in ref_to_id:
            ids.add(ref_to_id[ref])
    return ids


def score_case(case: dict, out: dict, k: int, ref_to_id: Dict[str, str]) -> dict:
    relevant = set(case["relevant"])
    results = out["results"][:k]

    retrieved, hit_ranks, found = [], [], set()
    for rank, item in enumerate(results, 1):
        ids = result_ids(item, ref_to_id)
        retrieved.append(item.get("id"))
        matched = ids & relevant
        if matched:
            hit_ranks.append(rank)
            found |= matched

    first = hit_ranks[0] if hit_ranks else None
    row = {
        "id": case["id"],
        "category": case["category"],
        "lang": case.get("lang", "?"),
        "query": case["query"],
        "intent": out["intent"],
        "top_score": round(float(out["top_score"]), 4),
        "gated": out["gated"],
        "confidence": None if out["confidence"] is None else round(float(out["confidence"]), 4),
        "confidence_source": out["confidence_source"],
        "retrieved": retrieved,
        "first_hit_rank": first,
    }
    if relevant:
        row.update({
            "hit@1": int(first is not None and first <= 1),
            "hit@3": int(first is not None and first <= 3),
            "hit@k": int(first is not None),
            "mrr": 1.0 / first if first else 0.0,
            # Deduped hadith results also match the narrations merged into them, so
            # found can exceed k; cap at 1.
            "recall@k": min(1.0, len(found) / min(len(relevant), k)),
        })
    if case.get("required"):
        missing = sorted(set(case["required"]) - found)
        row["complete@k"] = int(not missing)
        row["missing_required"] = missing
    return row


# ═══════════════════════════════════════════════════════════════════════
# REPORTING
# ═══════════════════════════════════════════════════════════════════════
METRICS = ("hit@1", "hit@3", "hit@k", "mrr", "recall@k")


def aggregate(rows: List[dict]) -> dict:
    pos = [r for r in rows if "hit@k" in r]
    neg = [r for r in rows if "hit@k" not in r]

    def mean(rs, key):
        return round(sum(r[key] for r in rs) / len(rs), 4) if rs else None

    def summarize(rs):
        return {"n": len(rs), **{m: mean(rs, m) for m in METRICS}}

    by_cat, by_lang = defaultdict(list), defaultdict(list)
    for r in pos:
        by_cat[r["category"]].append(r)
        by_lang[r["lang"]].append(r)

    req = [r for r in pos if "complete@k" in r]
    return {
        "overall": summarize(pos),
        "required": {"n": len(req), "complete@k": mean(req, "complete@k")},
        "by_category": {c: summarize(rs) for c, rs in sorted(by_cat.items())},
        "by_lang": {lang: summarize(rs) for lang, rs in sorted(by_lang.items())},
        "gate": _gate_summary(pos, neg),
    }


def _gate_summary(pos: List[dict], neg: List[dict]) -> Optional[dict]:
    if any(r["gated"] is None for r in pos + neg):
        return None
    rate = lambda flags: round(sum(flags) / len(flags), 4) if flags else None  # noqa: E731
    return {
        # Positives refused even though retrieval found the answer.
        "false_reject_rate": rate([int(r["gated"]) for r in pos if r["hit@k"]]),
        # Positives refused regardless of retrieval outcome (what the user experiences).
        "pos_reject_rate": rate([int(r["gated"]) for r in pos]),
        # Out-of-scope queries the gate lets through to the LLM.
        "negatives": len(neg),
        "false_accept_rate": rate([int(not r["gated"]) for r in neg]),
        "pos_confidence_median": _median([r["confidence"] for r in pos]),
        "neg_confidence_median": _median([r["confidence"] for r in neg]),
    }


def _median(xs: List[float]) -> Optional[float]:
    if not xs:
        return None
    s = sorted(xs)
    mid = len(s) // 2
    return round(s[mid] if len(s) % 2 else (s[mid - 1] + s[mid]) / 2, 4)


def _fmt(v) -> str:
    return "   —  " if v is None else f"{v:6.3f}"


def print_report(summary: dict, k: int) -> None:
    header = f"{'':22}{'n':>4}  " + "  ".join(f"{m.replace('@k', f'@{k}'):>8}" for m in METRICS)
    print(header)
    print("─" * len(header))

    def line(label, s):
        print(f"{label:22}{s['n']:>4}  " + "  ".join(f"{_fmt(s[m]):>8}" for m in METRICS))

    line("OVERALL", summary["overall"])
    print()
    for cat, s in summary["by_category"].items():
        line(f"  {cat}", s)
    print()
    for lang, s in summary["by_lang"].items():
        line(f"  lang={lang}", s)
    req = summary.get("required") or {}
    if req.get("n"):
        print(f"\n  complete@{k} (all must-include sources retrieved): {_fmt(req['complete@k'])}  (n={req['n']})")

    g = summary["gate"]
    if g is None:
        print("\nConfidence gate: not evaluated in hybrid mode")
        return
    print("\nConfidence gate")
    print(f"  false-reject rate (answer retrieved but gated): {_fmt(g['false_reject_rate'])}")
    print(f"  reject rate on all positives:                  {_fmt(g['pos_reject_rate'])}")
    print(f"  false-accept rate on {g['negatives']} negatives:          {_fmt(g['false_accept_rate'])}")
    print(f"  median confidence  positives={_fmt(g['pos_confidence_median'])}  "
          f"negatives={_fmt(g['neg_confidence_median'])}")


def print_failures(rows: List[dict], k: int) -> None:
    misses = [r for r in rows if r.get("hit@k") == 0]
    incomplete = [r for r in rows if r.get("hit@k") == 1 and r.get("complete@k") == 0]
    leaks = [r for r in rows if "hit@k" not in r and r["gated"] is False]
    refused = [r for r in rows if "hit@k" in r and r["gated"]]
    if misses:
        print(f"\nMisses (no relevant result in top {k}):")
        for r in misses:
            print(f"  {r['id']:<14} [{r['intent']}] {r['query'][:60]}")
            print(f"  {'':14} got: {', '.join(map(str, r['retrieved'][:5]))}")
    if incomplete:
        print(f"\nHits missing a must-include source:")
        for r in incomplete:
            print(f"  {r['id']:<14} {r['query'][:50]}  missing: {', '.join(r['missing_required'])}")
    for title, group in (("Negatives that pass the gate", leaks), ("Positives refused by the gate", refused)):
        if group:
            print(f"\n{title}:")
            for r in group:
                print(f"  {r['id']:<14} {r['confidence_source']}={r['confidence']:.3f}  "
                      f"{r['query'][:50]}  → {', '.join(map(str, r['retrieved'][:3]))}")


def print_comparison(current: dict, prev_path: Path) -> None:
    prev = json.loads(prev_path.read_text(encoding="utf-8"))
    print(f"\nΔ vs {prev_path.name} ({prev.get('timestamp', '?')}, commit {prev.get('git_commit', '?')})")

    def delta_line(label, cur, old):
        cells = []
        for m in METRICS:
            a, b = cur.get(m), (old or {}).get(m)
            cells.append("    —   " if a is None or b is None else f"{a - b:+8.3f}")
        print(f"{label:22}      " + "  ".join(cells))

    delta_line("OVERALL", current["summary"]["overall"], prev["summary"]["overall"])
    cur_req, old_req = current["summary"].get("required") or {}, prev["summary"].get("required") or {}
    if cur_req.get("complete@k") is not None and old_req.get("complete@k") is not None:
        print(f"{'  complete@k':22}      {cur_req['complete@k'] - old_req['complete@k']:+8.3f}")
    for cat, s in current["summary"]["by_category"].items():
        delta_line(f"  {cat}", s, prev["summary"]["by_category"].get(cat))

    old_rows = {r["id"]: r for r in prev["cases"]}
    fixed, broke = [], []
    for r in current["cases"]:
        o = old_rows.get(r["id"])
        if not o or "hit@k" not in r or "hit@k" not in o:
            continue
        if r["hit@k"] > o["hit@k"]:
            fixed.append(r["id"])
        elif r["hit@k"] < o["hit@k"]:
            broke.append(r["id"])
    print(f"\n  newly hit:  {', '.join(fixed) or '—'}")
    print(f"  newly lost: {', '.join(broke) or '—'}")


def _git_commit() -> str:
    try:
        sha = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"], cwd=ROOT, text=True).strip()
        dirty = subprocess.call(["git", "diff", "--quiet"], cwd=ROOT) != 0
        return sha + ("-dirty" if dirty else "")
    except Exception:
        return "unknown"


# ═══════════════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════════════
async def main_async(args: argparse.Namespace) -> None:
    from app.state import state

    only = set(args.only.split(",")) if args.only else None
    cases = load_golden(Path(args.golden), only)
    init_state(args.llm)
    validate_labels(cases, state.dataset)

    ref_to_id = {
        item["reference"]: item["id"]
        for item in state.dataset
        if item.get("type") == "hadith" and item.get("reference")
    }

    rows, t0 = [], time.perf_counter()
    for i, case in enumerate(cases, 1):
        out = await retrieve(case, args.mode, args.k)
        rows.append(score_case(case, out, args.k, ref_to_id))
        if sys.stdout.isatty():
            print(f"\r  {i}/{len(cases)} evaluated", end="", flush=True)
    elapsed = time.perf_counter() - t0
    print(f"\r  {len(cases)} cases in {elapsed:.1f}s ({elapsed / max(len(cases), 1):.2f}s/query)\n")

    summary = aggregate(rows)
    print_report(summary, args.k)
    if args.show_failures:
        print_failures(rows, args.k)

    report = {
        "timestamp": datetime.now().isoformat(timespec="seconds"),
        "git_commit": _git_commit(),
        "mode": args.mode,
        "k": args.k,
        "golden": str(Path(args.golden).resolve().relative_to(ROOT)),
        "config": {
            key: getattr(cfg, key) for key in (
                "EMBED_MODEL", "TOP_K_SEARCH", "TOP_K_RETURN", "USE_RERANKER", "RERANK_MODEL",
                "HADITH_BOOST", "QURAN_BOOST", "FATWA_BOOST", "RERANK_MIN_SCORE", "DENSE_MIN_SCORE",
                "FAST_QUERY_REWRITE", "ENABLE_HADITH_DEDUPLICATION",
            )
        },
        "summary": summary,
        "cases": rows,
    }

    if args.compare:
        print_comparison(report, Path(args.compare))

    out_path = Path(args.out) if args.out else RESULTS_DIR / f"{datetime.now():%Y%m%d-%H%M%S}-{args.mode}.json"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\nReport written to {out_path.relative_to(ROOT) if out_path.is_relative_to(ROOT) else out_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate QModel retrieval against a golden set.")
    parser.add_argument("--golden", default=str(DEFAULT_GOLDEN), help="Golden set JSONL path")
    parser.add_argument("--mode", choices=("pipeline", "hybrid"), default="pipeline")
    parser.add_argument("--k", type=int, default=cfg.TOP_K_RETURN,
                        help=f"Results to score (default TOP_K_RETURN={cfg.TOP_K_RETURN}, what the LLM sees)")
    parser.add_argument("--only", help="Comma-separated categories to run (e.g. fatwa,negative)")
    parser.add_argument("--llm", action="store_true", help="Load the LLM backend for query rewriting")
    parser.add_argument("--show-failures", action="store_true", help="List misses and leaked negatives")
    parser.add_argument("--compare", help="Previous report JSON to diff against")
    parser.add_argument("--out", help="Where to write the JSON report (default eval/results/<timestamp>.json)")
    parser.add_argument("--verbose", action="store_true", help="Show app logging")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO if args.verbose else logging.WARNING,
                        format="%(levelname)s %(name)s: %(message)s")
    asyncio.run(main_async(args))


if __name__ == "__main__":
    main()