#!/usr/bin/env python3
"""Score captured retrieval rows against independent relevance judgments."""

from __future__ import annotations

import argparse
import json
import statistics
from pathlib import Path
from typing import Any


def score(results_path: Path, judgments_path: Path) -> dict[str, Any]:
    lines = [line for line in results_path.read_text(encoding="utf-8").splitlines() if line]
    if len(lines) < 2:
        raise ValueError("Retrieval results must contain a manifest and at least one query")
    manifest = json.loads(lines[0])
    rows = [json.loads(line) for line in lines[1:]]
    question_ids = [row.get("question", {}).get("id") for row in rows]
    if any(not isinstance(question_id, str) for question_id in question_ids):
        raise ValueError("Every retrieval row must identify its question")
    if len(question_ids) != len(set(question_ids)):
        raise ValueError("Retrieval results contain duplicate question IDs")
    if len(rows) != manifest.get("question_count"):
        raise ValueError("Captured query count does not match the run manifest")

    judgments = json.loads(judgments_path.read_text(encoding="utf-8"))
    relevant_by_question = judgments.get("fully_answerable_positive_questions")
    if not isinstance(relevant_by_question, dict) or not relevant_by_question:
        raise ValueError("Independent judgments contain no answerable positive questions")
    result_by_question = {row["question"]["id"]: row for row in rows}
    missing = sorted(set(relevant_by_question) - result_by_question.keys())
    if missing:
        raise ValueError(f"Retrieval run omitted judged questions: {', '.join(missing)}")

    scored: list[dict[str, Any]] = []
    for question_id, relevant_ids in relevant_by_question.items():
        row = result_by_question[question_id]
        if row.get("is_error"):
            returned_ids: list[str] = []
        else:
            returned_ids = row.get("returned_node_ids", [])
        relevant = set(relevant_ids)
        hits = relevant & set(returned_ids[:5])
        scored.append(
            {
                "question_id": question_id,
                "relevant_node_ids": sorted(relevant),
                "returned_node_ids": returned_ids[:5],
                "hit_node_ids": sorted(hits),
                "recall_at_5": len(hits) / len(relevant),
                "is_error": bool(row.get("is_error")),
            }
        )

    revisions = {
        page.get("canonical_revision")
        for row in rows
        for page in row.get("response", [])
        if page.get("canonical_revision") is not None
    }
    if len(revisions) > 1:
        raise ValueError("Retrieval responses came from more than one canonical revision")
    hit_count = sum(bool(row["hit_node_ids"]) for row in scored)
    query_times = [float(row["seconds"]) for row in rows]
    return {
        "schema_version": 1,
        "questions_sha256": manifest.get("questions_sha256"),
        "question_count": len(rows),
        "judged_answerable_questions": len(scored),
        "mean_conditional_recall_at_5": statistics.mean(row["recall_at_5"] for row in scored),
        "conditional_hit_rate": hit_count / len(scored),
        "queries_with_any_rows": sum(bool(row.get("returned_node_ids")) for row in rows),
        "mean_observed_seconds": statistics.mean(query_times),
        "median_observed_seconds": statistics.median(query_times),
        "canonical_revision": next(iter(revisions), None),
        "partial_ingestion_questions": sorted(judgments.get("partial_ingestion", [])),
        "missing_accepted_material_questions": sorted(
            judgments.get("missing_accepted_material", [])
        ),
        "pagination_blocked_queries": sum(bool(row.get("pagination_blocked")) for row in rows),
        "judgments_method": judgments.get("method", "not supplied"),
        "per_question": scored,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--results", type=Path, required=True)
    parser.add_argument("--judgments", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    result = score(args.results, args.judgments)
    with args.output.open("x", encoding="utf-8") as output:
        output.write(json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n")
    print(
        json.dumps(
            {
                "questions": result["question_count"],
                "judged_answerable": result["judged_answerable_questions"],
                "conditional_recall_at_5": result["mean_conditional_recall_at_5"],
                "conditional_hit_rate": result["conditional_hit_rate"],
                "output": str(args.output),
            }
        )
    )


if __name__ == "__main__":
    main()