"""
test_rag.py — run with `pytest -q` (offline, deterministic).

These tests are the LAB'S SUCCESS CRITERIA expressed as code. Each asserts a
concept from Knowledge Module 02. If they pass, you have a working hybrid RAG.
"""
from __future__ import annotations

from data import DOCUMENTS, EVAL
from rag import (BM25, Embedder, LocalHybridIndex, RagPipeline, chunk_text,
                 cosine, rrf_fuse, tokenize)


def build() -> RagPipeline:
    pipe = RagPipeline()
    pipe.ingest(DOCUMENTS)
    return pipe


# --- unit: the building blocks ---------------------------------------------

def test_chunking_overlaps_and_covers():
    text = " ".join(f"Sentence number {i} has content." for i in range(40))
    chunks = chunk_text(text, max_tokens=30, overlap=8)
    assert len(chunks) > 1                      # it actually split
    assert all(len(tokenize(c)) <= 60 for c in chunks)  # roughly bounded


def test_bm25_ranks_exact_identifier_first():
    bm = BM25()
    bm.index(["a", "b"], [tokenize("the swift wire transfer costs 50 aed"),
                          tokenize("local transfers are free of charge")])
    top = bm.search("swift transfer cost", top_k=2)
    assert top[0][0] == "a"                     # exact keyword doc wins


def test_vector_matches_paraphrase_via_synonyms():
    emb = Embedder()
    docs = [tokenize("a card transaction can be refused for insufficient balance"),
            tokenize("monthly statements are generated on the first day")]
    emb.fit(docs)
    q = emb.embed("why was my payment declined")   # 'declined' ~ 'refused' synonym
    sims = [cosine(q, emb.embed(" ".join(d))) for d in docs]
    assert sims[0] > sims[1]                     # semantic match beats unrelated doc


def test_rrf_rewards_agreement():
    a = [("x", 9), ("y", 8), ("z", 7)]
    b = [("y", 5), ("x", 4)]
    fused = dict(rrf_fuse([a, b]))
    assert fused["y"] > fused["z"]               # y is high in both lists -> wins over z


# --- behavior: grounding, abstention, security trimming --------------------

def test_answer_is_cited():
    pipe = build()
    r = pipe.ask("How much is a paper statement?", user_groups=["retail"])
    assert not r["abstained"]
    assert r["citations"] and r["citations"][0].startswith("[statements:")
    assert "10" in r["answer"]


def test_abstains_when_out_of_corpus():
    pipe = build()
    r = pipe.ask("What is the bank's stock price today?", user_groups=["retail"])
    assert r["abstained"] is True
    assert r["grounded"] is True                 # abstaining is grounded behavior


def test_security_trimming_hides_staff_doc():
    pipe = build()
    # The override code lives in a STAFF-only doc. A retail user must NOT see it —
    # the secret never appears, and the staff doc is never even retrieved. (The
    # user may still get a legitimate, non-secret answer about fee waivers.)
    r = pipe.ask("How can a branch waive my account fee with an override code?",
                 user_groups=["retail"])
    assert "OVR-204" not in r["answer"]
    assert all("internal-override" not in cid for cid, _ in r["retrieved"])
    # A staff user, by contrast, CAN retrieve it.
    r2 = pipe.ask("fee override code", user_groups=["staff"])
    assert any("internal-override" in cid for cid, _ in r2["retrieved"])


def test_language_filter():
    pipe = build()
    r = pipe.ask("رسوم الحساب الجاري", user_groups=["retail"], language="ar")
    assert all("-ar" in cid or cid.startswith("fees-current-account-ar")
               for cid, _ in r["retrieved"]) or r["abstained"]


# --- the thesis: hybrid >= its parts on recall@5 ---------------------------

def _recall_at_5(pipe: RagPipeline, mode: str) -> float:
    hits = 0
    for case in EVAL:
        if case.expected_doc == "internal-override":
            continue  # excluded: it's a security-trimming case, not a recall case
        r = pipe.ask(case.question, mode=mode, top_k=5, user_groups=["retail"])
        docs = {cid.split("#")[0] for cid, _ in r["retrieved"]}
        hits += case.expected_doc in docs
    return hits / (len(EVAL) - 1)


def test_hybrid_beats_or_matches_components():
    pipe = build()
    vector = _recall_at_5(pipe, "vector")
    bm25 = _recall_at_5(pipe, "bm25")
    hybrid = _recall_at_5(pipe, "hybrid")
    hybrid_sem = _recall_at_5(pipe, "hybrid_semantic")
    # Hybrid should be at least as good as the better single component.
    assert hybrid >= max(vector, bm25) - 1e-9
    assert hybrid_sem >= 0.8                      # strong absolute recall