"""
run.py — Lab 01 demo. Runs fully offline.

    python run.py            # ask a few sample questions + show the A/B retrieval table
    python run.py "How long do I have to dispute a transaction?"   # ask your own

What you should SEE and be able to EXPLAIN in the interview:
  1. Grounded answers with citations, and ABSTENTION when the answer isn't in scope.
  2. SECURITY TRIMMING: a 'retail' user asking about the staff-only fee override
     gets an abstention, not the secret override code.
  3. The A/B table: hybrid (and hybrid+semantic) recall@5 >= vector or bm25 alone.
"""
from __future__ import annotations
import sys

from data import DOCUMENTS, EVAL
from rag import RagPipeline


def demo_questions(pipe: RagPipeline) -> None:
    print("=" * 72)
    print("SAMPLE ANSWERS (mode=hybrid_semantic, user=retail)")
    print("=" * 72)
    for q in [
        "How much does an international SWIFT wire transfer cost?",
        "My card payment was refused — why might that happen?",
        "Can I get my account fee waived by talking to a branch?",  # staff-only -> abstain
    ]:
        r = pipe.ask(q, user_groups=["retail"])
        flag = "  (ABSTAINED)" if r["abstained"] else ""
        print(f"\nQ: {q}")
        print(f"A: {r['answer']}{flag}")
        print(f"   retrieved: {r['retrieved']}")


def ab_table(pipe: RagPipeline) -> None:
    print("\n" + "=" * 72)
    print("A/B: retrieval mode vs recall@5 and answer-correctness (20-question eval)")
    print("=" * 72)
    modes = ["bm25", "vector", "hybrid", "hybrid_semantic"]
    print(f"{'mode':<18}{'recall@5':>10}{'correct':>10}")
    for mode in modes:
        hits = correct = 0
        for case in EVAL:
            groups = ["retail"]  # all eval users are retail customers
            r = pipe.ask(case.question, mode=mode, top_k=5, user_groups=groups)
            retrieved_docs = {cid.split("#")[0] for cid, _ in r["retrieved"]}
            # The override question SHOULD NOT retrieve the staff doc for a retail user:
            if case.expected_doc == "internal-override":
                if r["abstained"]:
                    hits += 1
                    correct += 1
                continue
            if case.expected_doc in retrieved_docs:
                hits += 1
            if (not r["abstained"]) and case.key_fact and case.key_fact.lower() in r["answer"].lower():
                correct += 1
        n = len(EVAL)
        print(f"{mode:<18}{hits/n:>10.2f}{correct/n:>10.2f}")
    print("\nNote: hybrid fuses BM25 (exact identifiers like '5.99', 'SWIFT') with the")
    print("vector retriever (paraphrases like 'refused' ~ 'declined'), and the semantic")
    print("reranker reorders the fused set — so it matches or beats either component.")


def main() -> None:
    pipe = RagPipeline()
    pipe.ingest(DOCUMENTS)
    if len(sys.argv) > 1:
        q = " ".join(sys.argv[1:])
        r = pipe.ask(q, user_groups=["retail"])
        print(f"Q: {q}\nA: {r['answer']}")
        print(f"mode={r['mode']}  retrieved={r['retrieved']}  abstained={r['abstained']}")
        return
    demo_questions(pipe)
    ab_table(pipe)


if __name__ == "__main__":
    main()