"""Measure the retriever against the generated dataset.

Scores every retrieval mode on hit rate and MRR.

Free to run, because embedding is local and no API calls are made, so we can
re-run it after every retriever change.

    uv run python -m eval.run_eval
    uv run python -m eval.run_eval --retriever hybrid --retriever bm25
"""

import argparse
import json
import statistics
import sys

from llama_index.core.evaluation import HitRate, MRR

from rag.config import CHUNKS_PATH, EVAL_DATASET_PATH, RETRIEVAL_MODE, RETRIEVAL_TOP_K
from rag.index import RETRIEVAL_MODES, load_retriever, make_embed_model
from rag.types import EvalRow

# What counts as correct: only the chunk the question came from, or any chunk of
# the same section. We score both, because a neighbouring chunk often answers the
# question just as well and a citation points at the section anyway.
LEVELS = ("strict", "source")
METRICS = tuple(f"{level}_{metric}" for level in LEVELS for metric in ("found", "rank"))

# One score per question, per metric, averaged once every question is scored.
Scores = dict[str, list[float]]

HIT_RATE, RECIPROCAL_RANK = HitRate(), MRR()


def rank_all(retriever, rows: list[EvalRow]) -> list[list[str]]:
    return [[node.node_id for node in retriever.retrieve(row["question"])] for row in rows]


def same_source_ids() -> dict[str, list[str]]:
    by_url: dict[str, list[str]] = {}
    ids: dict[str, str] = {}
    with CHUNKS_PATH.open(encoding="utf-8") as handle:
        for line in handle:
            chunk = json.loads(line)
            url = chunk["metadata"]["url"]
            ids[chunk["id"]] = url
            by_url.setdefault(url, []).append(chunk["id"])
    return {chunk_id: by_url[url] for chunk_id, url in ids.items()}


def score_questions(
    rows: list[EvalRow], rankings: list[list[str]], expand: dict[str, list[str]]
) -> Scores:
    scores: Scores = {metric: [] for metric in METRICS}
    for row, ranked in zip(rows, rankings):
        retrieved = ranked[:RETRIEVAL_TOP_K]
        expected = {"strict": [row["chunk_id"]], "source": expand[row["chunk_id"]]}
        for level in LEVELS:
            ids = expected[level]
            scores[f"{level}_found"].append(
                HIT_RATE.compute(expected_ids=ids, retrieved_ids=retrieved).score
            )
            scores[f"{level}_rank"].append(
                RECIPROCAL_RANK.compute(expected_ids=ids, retrieved_ids=retrieved).score
            )
    return scores


def report(scores: dict[str, Scores]) -> None:
    print(f"\n{'':12}{'exact':>18}{'section':>20}")
    print(f"{'':12}{'hit_rate':>9}{'mrr':>9}{'hit_rate':>11}{'mrr':>9}")
    for mode, metrics in scores.items():
        marker = "*" if mode == RETRIEVAL_MODE else " "
        cells = "".join(
            f"{statistics.fmean(metrics[m]):>9.2f}" + ("  " if m == "strict_rank" else "")
            for m in METRICS
        )
        print(f"{marker} {mode:10}{cells}")
    print(f"\n* scored on the top {RETRIEVAL_TOP_K} results")


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--retriever", action="append", choices=RETRIEVAL_MODES,
        help="modes to measure; repeatable, defaults to all",
    )
    parser.add_argument("--rerank-model", help="override the cross-encoder for rerank mode")
    args = parser.parse_args(argv[1:])
    modes = args.retriever or list(RETRIEVAL_MODES)

    if not EVAL_DATASET_PATH.exists():
        print("No dataset. Run: uv run python -m eval.build_dataset", file=sys.stderr)
        return 2
    rows = json.loads(EVAL_DATASET_PATH.read_text(encoding="utf-8"))

    # Same embedding model and query instruction the app uses
    print(f"{len(rows)} questions | query instruction: {make_embed_model().query_instruction!r}")

    expand = same_source_ids()
    scores: dict[str, Scores] = {}
    for mode in modes:
        retriever = load_retriever(mode, RETRIEVAL_TOP_K, None, args.rerank_model)
        scores[mode] = score_questions(rows, rank_all(retriever, rows), expand)
        print(f"  scored {mode}", flush=True)

    report(scores)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))