#!/usr/bin/env python
"""End-to-end check on a real document: the NIST AI Risk Management Framework 1.0.

Downloads NIST AI 100-1 (January 2023, 48 pages, a US-government work in the public
domain), indexes it, and asks ten questions with known outcomes: six the document
answers, three near misses (on-topic, but the fact is not in it) and one off-topic.
Each question runs under three setups so the effect of each abstention gate shows:
no gates, the retrieval-similarity gate, and the similarity gate plus the reader.

This is a spot check on one document, not a benchmark; see scripts/benchmark.py
for calibrated numbers.

``--retrieval`` instead measures retrieval inside the document with and without
contextual chunk headers, on 22 questions graded by verbatim evidence spans
(data/nist_rmf/eval.jsonl).

Usage (needs the embed, reader and pdf extras; downloads ~200 MB of models once):
  python scripts/real_pdf_check.py
  python scripts/real_pdf_check.py --retrieval   # embed and pdf extras only
"""
from __future__ import annotations

import argparse
import sys
import tempfile
import time
from pathlib import Path

import httpx

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

from lumen_rag import config
from lumen_rag.embeddings import FastEmbedEmbedder, HashingEmbedder
from lumen_rag.engine import RagEngine
from lumen_rag.eval.evidence import evaluate_evidence, load_evidence_cases
from lumen_rag.ingestion.loaders import load_file
from lumen_rag.ingestion.pipeline import ingest_documents
from lumen_rag.retrieval import Retriever
from lumen_rag.store import VectorStore

PDF_URL = "https://nvlpubs.nist.gov/nistpubs/ai/NIST.AI.100-1.pdf"
READER = "onnx-community/roberta-base-squad2-ONNX"

# (kind, question): "answerable", "near-miss" or "off-topic".
QUESTIONS = [
    ("answerable", "Who was the Secretary of Commerce when the AI RMF was published?"),
    ("answerable", "Which function is designed to be cross-cutting across the other three?"),
    ("answerable", "What does TEVV stand for?"),
    ("answerable", "Where can comments on the AI RMF Playbook be sent?"),
    ("answerable", "When is a formal review of the framework expected to take place?"),
    ("answerable", "What are the characteristics of trustworthy AI systems?"),
    ("near-miss", "When will version 2.0 of the AI RMF be released?"),
    ("near-miss", "How much does AI RMF certification cost?"),
    ("near-miss", "What fine does an organization pay for not following the AI RMF?"),
    ("off-topic", "What is the boiling point of water at sea level?"),
]

SETUPS = [
    ("no gates", None, ""),
    ("similarity >= 0.66", 0.66, ""),
    ("similarity >= 0.66, then reader >= -3.5", 0.66, READER),
]


def download(tmp: str) -> Path:
    pdf = Path(tmp) / "NIST.AI.100-1.pdf"
    print(f"Downloading {PDF_URL}")
    # nvlpubs.nist.gov refuses Python's default user agent.
    resp = httpx.get(PDF_URL, follow_redirects=True, timeout=60, headers={"User-Agent": "lumen-rag/0.1"})
    resp.raise_for_status()
    pdf.write_bytes(resp.content)
    return pdf


def retrieval(pdf: Path) -> None:
    """hit@5 and MRR inside the document, with and without contextual chunk headers."""
    cases = load_evidence_cases(ROOT / "data" / "nist_rmf" / "eval.jsonl")
    doc = load_file(pdf)
    print(f"{len(cases)} questions graded by evidence spans, k=5, 120-word chunks\n")
    print("| Embedder | Mode | hit@5 | MRR | hit@5 + headers | MRR + headers |")
    print("|---|---|---|---|---|---|")
    for name, embedder in (("hashing", HashingEmbedder()), ("bge-small", FastEmbedEmbedder())):
        stores = {}
        for ctx in (False, True):
            stores[ctx] = VectorStore(embedder.dim)
            ingest_documents([doc], store=stores[ctx], embedder=embedder, contextual=ctx)
        for mode in ("vector", "bm25", "hybrid"):
            plain, headed = (
                evaluate_evidence(Retriever(stores[ctx], embedder, reranker=None), cases, k=5, mode=mode)
                for ctx in (False, True)
            )
            print(
                f"| {name} | {mode} | {plain['hit@k']:.2f} | {plain['mrr']:.2f} "
                f"| {headed['hit@k']:.2f} | {headed['mrr']:.2f} |"
            )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--retrieval", action="store_true", help="compare retrieval with and without headers")
    args = parser.parse_args()
    config.settings.embedder = "fastembed"
    config.settings.reader_threshold = -3.5
    with tempfile.TemporaryDirectory() as tmp:
        pdf = download(tmp)
        if args.retrieval:
            retrieval(pdf)
            return
        engine = RagEngine()
        doc = load_file(pdf)
        total = engine.add_documents([doc])
        print(f"Indexed {pdf.name}: {total} chunks\n")

        summary = []
        for name, threshold, reader in SETUPS:
            config.settings.abstain_threshold = threshold
            config.settings.reader_model = reader
            right = 0
            print(f"=== {name}")
            for kind, question in QUESTIONS:
                start = time.perf_counter()
                result = engine.query(question, k=5, mode="vector")
                ms = (time.perf_counter() - start) * 1000
                ok = result.abstained == (kind != "answerable")
                right += ok
                sim = "-" if result.confidence is None else f"{result.confidence:.3f}"
                read = "-" if result.answerability is None else f"{result.answerability:+.1f}"
                verdict = "refused " if result.abstained else "answered"
                print(f"  [{'ok' if ok else 'XX'}] {kind:10} {verdict} sim={sim} reader={read} {ms:5.0f} ms  {question}")
            summary.append((name, right))
            print()

        print("| Setup | Right outcome |")
        print("|---|---|")
        for name, right in summary:
            print(f"| {name} | {right}/{len(QUESTIONS)} |")


if __name__ == "__main__":
    main()