"""Page-level miss dump using this clone's Dense/BM25 retrievers.

    python -m eval.trace --queries eval/queries.jsonl
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from eval.run_eval import load_queries, p_at_1


def classify_top1_source(top1, dense_pages, bm25_pages) -> str:
    if top1 is None:
        return "neither"
    in_d, in_b = top1 in dense_pages, top1 in bm25_pages
    if in_d and in_b:
        return "both"
    if in_d:
        return "dense_only"
    if in_b:
        return "bm25_only"
    return "neither"


def first_ranks(pages: list[tuple[str, str]]) -> dict[str, int]:
    out: dict[str, int] = {}
    for i, (bid, pid) in enumerate(pages, start=1):
        key = f"{bid}/{pid}"
        if key not in out:
            out[key] = i
    return out


def render_misses(rows: list[dict]) -> str:
    misses = [r for r in rows if not r.get("hit")]
    if not misses:
        return "# Eval misses\n\n_(none)_\n"
    lines = [f"# Eval misses ({len(misses)})", ""]
    for row in misses:
        lines += [
            f"## {row['query']}",
            "",
            f"- expected: `{row['expected']}`",
            f"- top-5: `{row['ranked_pages'][:5]}`",
            f"- top1_source: `{row['top1_source']}`",
            "",
        ]
    return "\n".join(lines)


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--queries", type=Path, default=Path("eval/queries.jsonl"))
    p.add_argument("--trace-out", type=Path, default=Path("eval/results/trace.jsonl"))
    p.add_argument("--misses-out", type=Path, default=Path("eval/results/misses.md"))
    args = p.parse_args()
    if not args.queries.exists():
        print(f"queries file not found: {args.queries}")
        return 1
    try:
        import os
        import chromadb
        from chromadb.config import Settings
        from rag.embeddings import embedding_function
        from rag.bm25 import BM25Index
        from eval.retrievers import BM25Retriever, DenseRetriever
        product = os.environ.get("PRODUCT_NAME", "hvm")
        root = Path(__file__).resolve().parent.parent
        col = chromadb.PersistentClient(
            path=str(root / "chroma"),
            settings=Settings(anonymized_telemetry=False),
        ).get_collection(f"{product}_docs", embedding_function=embedding_function())
        bm = BM25Index(str(root / "bm25" / f"{product}_docs.db"))
        dense_r, bm25_r = DenseRetriever(col), BM25Retriever(bm)
    except Exception as e:
        args.trace_out.parent.mkdir(parents=True, exist_ok=True)
        args.trace_out.write_text("")
        args.misses_out.write_text("# Eval misses\n\nno index\n")
        print(f"no index ({e}); wrote empty trace")
        return 0

    rows = []
    for q in load_queries(args.queries):
        expected = [(e["bundle_id"], e["page_id"]) for e in q["expected"]]
        dense_pages = dense_r.retrieve(q["query"], k=50)
        bm25_pages = bm25_r.retrieve(q["query"], k=50)
        ranked = bm25_pages or dense_pages  # HVM default retrieval is BM25-first
        top1 = ranked[0] if ranked else None
        p1 = p_at_1(ranked, expected)
        rows.append({
            "query": q["query"],
            "expected": [list(p) for p in expected],
            "hit": bool(p1),
            "p_at_1": int(p1),
            "dense_rank": first_ranks(dense_pages),
            "bm25_rank": first_ranks(bm25_pages),
            "top1": list(top1) if top1 else None,
            "top1_source": classify_top1_source(top1, set(dense_pages), set(bm25_pages)),
            "ranked_pages": [list(p) for p in ranked],
        })
    args.trace_out.parent.mkdir(parents=True, exist_ok=True)
    with open(args.trace_out, "w") as fh:
        for row in rows:
            fh.write(json.dumps(row) + "\n")
    args.misses_out.write_text(render_misses(rows))
    print(f"wrote {args.trace_out} ({len(rows)} queries)")
    print(f"wrote {args.misses_out}")
    return 0


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