"""Persona tests for the shared secure retrieval function (PB-RAG-04-02).

Two layers:
- Offline tests (always run): fake Microsoft Graph and fake Azure AI Search, so CI needs no tenant.
- Live tests (RUN_LIVE_TRIMMING_TESTS=1): the same personas against the seeded sandbox index.

Run: pytest -q snippets/test_trimming.py
"""

from __future__ import annotations

import logging
import os
import re
import sys
from pathlib import Path
from typing import Any

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent))
import secure_search as ss  # noqa: E402

# Synthetic ids. They are not real tenant objects.
G_ALL = "00000000-0000-4000-8000-0000000000a1"
G_FIN = "00000000-0000-4000-8000-0000000000a2"
G_LEGAL = "00000000-0000-4000-8000-0000000000a3"
G_LEGAL_UK = "00000000-0000-4000-8000-0000000000a4"  # member of G_LEGAL (nested)
U_ANNA = "00000000-0000-4000-8000-0000000000b1"  # finance
U_BEN = "00000000-0000-4000-8000-0000000000b2"  # legal-uk, so legal through nesting
U_CAROL = "00000000-0000-4000-8000-0000000000b3"  # all staff, owns a personal document
U_DAN = "00000000-0000-4000-8000-0000000000b4"  # more than 200 groups (token overage)
U_OUTAGE = "00000000-0000-4000-8000-0000000000b5"
U_THROTTLED = "00000000-0000-4000-8000-0000000000b6"

MANY = [f"00000000-0000-4000-8000-{i:012x}" for i in range(0x1000, 0x1000 + 300)]

DOCS = {
    "doc-public": {"group_ids": [G_ALL], "user_ids": [], "category": "general"},
    "doc-finance": {"group_ids": [G_FIN], "user_ids": [], "category": "finance"},
    "doc-restricted": {"group_ids": [G_LEGAL], "user_ids": [], "category": "legal"},
    "doc-personal": {"group_ids": [], "user_ids": [U_CAROL], "category": "general"},
}


def _grp(gid: str) -> dict[str, str]:
    return {"@odata.type": ss.GROUP_TYPE, "id": gid}


# Graph returns transitive memberships, so Ben already gets G_LEGAL through G_LEGAL_UK.
MEMBERSHIP_PAGES: dict[str, list[list[dict[str, str]]]] = {
    U_ANNA: [[_grp(G_FIN), _grp(G_ALL)]],
    U_BEN: [[_grp(G_LEGAL_UK), _grp(G_LEGAL), _grp(G_ALL),
             {"@odata.type": "#microsoft.graph.directoryRole", "id": "00000000-0000-4000-8000-0000000000c1"}]],
    U_CAROL: [[_grp(G_ALL)]],
    U_DAN: [[_grp(g) for g in MANY[:100]], [_grp(g) for g in MANY[100:200]],
            [_grp(g) for g in MANY[200:]] + [_grp(G_ALL)]],
}


class FakeResponse:
    def __init__(self, status: int, body: dict[str, Any] | None = None, headers: dict[str, str] | None = None):
        self.status_code = status
        self._body = body or {}
        self.headers = headers or {}

    def json(self) -> dict[str, Any]:
        return self._body


class FakeSession:
    """Plays Microsoft Graph and Azure AI Search. Evaluates only the filter shapes secure_search builds."""

    def __init__(self, search_status: int = 200):
        self.graph_calls: list[str] = []
        self.search_bodies: list[dict[str, Any]] = []
        self.search_status = search_status

    def get(self, url: str, headers: dict[str, str], timeout: float) -> FakeResponse:
        self.graph_calls.append(url)
        m = re.search(r"/users/([0-9a-f-]+)/transitiveMemberOf(?:\?page=(\d+))?$", url)
        assert m, url
        oid, page = m.group(1), int(m.group(2) or 0)
        if oid == U_OUTAGE:
            return FakeResponse(503)
        if oid == U_THROTTLED:
            return FakeResponse(429, headers={"Retry-After": "30"})
        pages = MEMBERSHIP_PAGES.get(oid, [[]])
        body: dict[str, Any] = {"value": pages[page]}
        if page + 1 < len(pages):
            body["@odata.nextLink"] = f"{ss.GRAPH_BASE}/users/{oid}/transitiveMemberOf?page={page + 1}"
        return FakeResponse(200, body)

    def post(self, url: str, headers: dict[str, str], json: dict[str, Any], timeout: float) -> FakeResponse:
        self.search_bodies.append(json)
        if self.search_status != 200:
            return FakeResponse(self.search_status)
        flt = json.get("filter") or ""
        um = re.search(r"u eq '([0-9a-f-]+)'", flt)
        if um is None:  # no security clause: behave like an untrimmed index
            return FakeResponse(200, {"value": [{"chunk_id": k} for k in DOCS]})
        user = um.group(1)
        m = re.search(r"search\.in\(g, '([^']*)', ','\)", flt)
        groups = set(m.group(1).split(",")) if m else set()
        category = re.search(r"\(category eq '((?:[^']|'')*)'\)", flt)
        hits = []
        for key, doc in DOCS.items():
            visible = user in doc["user_ids"] or bool(groups & set(doc["group_ids"]))
            if category and doc["category"] != category.group(1).replace("''", "'"):
                visible = False
            if visible:
                hits.append({"chunk_id": key})
        return FakeResponse(200, {"value": hits})


class Clock:
    def __init__(self) -> None:
        self.now = 1000.0

    def __call__(self) -> float:
        return self.now


def make(search_status: int = 200, ttl: int = 300, clock: Clock | None = None):
    settings = ss.Settings(search_endpoint="https://search.invalid", index_name="${INDEX_ALIAS}",
                           cache_ttl_seconds=ttl, allowed_filter_fields=frozenset({"category"}))
    session = FakeSession(search_status)
    cache = ss.GroupCache(ttl, clock or Clock())
    retriever = ss.SecureRetriever(settings, token_provider=lambda scope: "test-token", session=session, cache=cache)
    return retriever, session


def ids(hits: list[dict[str, Any]]) -> set[str]:
    return {h["chunk_id"] for h in hits}


PERSONAS = [
    ("finance_user", U_ANNA, {"doc-public", "doc-finance"}, {"doc-restricted", "doc-personal"}),
    ("nested_group_user", U_BEN, {"doc-public", "doc-restricted"}, {"doc-finance", "doc-personal"}),
    ("direct_user_grant", U_CAROL, {"doc-public", "doc-personal"}, {"doc-finance", "doc-restricted"}),
    ("overage_user", U_DAN, {"doc-public"}, {"doc-finance", "doc-restricted", "doc-personal"}),
    ("graph_outage", U_OUTAGE, set(), set(DOCS)),
    ("graph_throttled", U_THROTTLED, set(), set(DOCS)),
]


@pytest.mark.parametrize("vector", [None, [0.1, 0.2, 0.3]], ids=["keyword", "hybrid"])
@pytest.mark.parametrize("name,user,allowed,denied", PERSONAS, ids=[p[0] for p in PERSONAS])
def test_persona_sees_only_permitted_documents(name, user, allowed, denied, vector):
    retriever, session = make()
    got = ids(retriever.search("policy", user, vector=vector))
    assert not got & denied, f"{name} leaked {sorted(got & denied)}"
    assert allowed <= got, f"{name} over-blocked {sorted(allowed - got)}"
    for body in session.search_bodies:
        assert "filterOverride" not in str(body)
        if vector is not None:
            assert body["vectorFilterMode"] == "preFilter"


def test_graph_failure_never_reaches_the_index():
    retriever, session = make()
    assert retriever.search("policy", U_OUTAGE) == []
    assert session.search_bodies == []


def test_invalid_user_id_fails_closed_without_calls():
    retriever, session = make()
    assert retriever.search("policy", "anna@example.invalid") == []
    assert session.graph_calls == [] and session.search_bodies == []


def test_search_service_error_returns_nothing():
    retriever, _ = make(search_status=403)
    assert retriever.search("policy", U_ANNA) == []


def test_overage_user_pages_through_graph_and_uses_one_search_in_clause():
    retriever, session = make()
    retriever.search("policy", U_DAN)
    assert len(session.graph_calls) == 3
    flt = session.search_bodies[0]["filter"]
    assert flt.count("search.in(") == 1
    assert len(re.search(r"search\.in\(g, '([^']*)'", flt).group(1).split(",")) == 301


def test_non_group_objects_are_ignored():
    retriever, session = make()
    retriever.search("policy", U_BEN)
    assert "0000000000c1" not in session.search_bodies[0]["filter"]


def test_business_filter_is_anded_and_escaped():
    retriever, session = make()
    assert ids(retriever.search("q", U_ANNA, business_filters={"category": "finance"})) == {"doc-finance"}
    hostile = "x') or true or ('"
    got = ids(retriever.search("q", U_ANNA, business_filters={"category": hostile}))
    assert got == set()
    flt = session.search_bodies[-1]["filter"]
    assert flt.startswith("(user_ids/any(") and ") and ((category eq 'x'') or true or (''')" in flt
    retriever.search("q", U_ANNA, business_filters={"category": ["finance", hostile]})
    assert "search.in(category, 'finance|x'') or true or (''', '|')" in session.search_bodies[-1]["filter"]


@pytest.mark.parametrize("bad", [{"secret_field": "a"}, {"group_ids": G_FIN}, {"category": []}, {"category": ["a|b"]}])
def test_business_filter_rejects_unsafe_input(bad):
    retriever, _ = make()
    retriever.settings = ss.Settings(**{**retriever.settings.__dict__,
                                        "allowed_filter_fields": frozenset({"category", "group_ids"})})
    with pytest.raises(ValueError):
        retriever.search("q", U_ANNA, business_filters=bad)


def test_group_cache_expires_and_does_not_cache_failures():
    clock = Clock()
    retriever, session = make(ttl=300, clock=clock)
    retriever.search("q", U_ANNA)
    retriever.search("q", U_ANNA)
    assert len(session.graph_calls) == 1
    clock.now += 301
    retriever.search("q", U_ANNA)
    assert len(session.graph_calls) == 2
    retriever.search("q", U_OUTAGE)
    retriever.search("q", U_OUTAGE)
    assert sum(U_OUTAGE in c for c in session.graph_calls) == 2


def test_audit_log_has_hash_not_group_list(caplog):
    retriever, _ = make()
    with caplog.at_level(logging.INFO, logger="secure_search"):
        retriever.search("q", U_ANNA)
    text = caplog.text
    assert "filter_sha256=" in text and "groups=2" in text
    assert G_FIN not in text and G_ALL not in text


def test_suggest_is_trimmed():
    retriever, session = make()
    retriever.suggest("pol", U_CAROL, "sg")
    assert session.search_bodies[0]["filter"].startswith(f"(user_ids/any(u: u eq '{U_CAROL}')")


def test_suggest_fails_closed_on_graph_outage():
    retriever, session = make()
    assert retriever.suggest("pol", U_OUTAGE, "sg") == []
    assert session.search_bodies == []


# ---- Live tests against the seeded sandbox index -------------------------------------------------

LIVE = os.environ.get("RUN_LIVE_TRIMMING_TESTS") == "1"


@pytest.mark.skipif(not LIVE, reason="set RUN_LIVE_TRIMMING_TESTS=1 to run against the sandbox index")
@pytest.mark.parametrize("env_user,should_see", [
    ("LIVE_ALLOWED_USER_OID", True),
    ("LIVE_NESTED_USER_OID", True),
    ("LIVE_DENIED_USER_OID", False),
])
def test_live_restricted_document(env_user, should_see):
    retriever = ss.SecureRetriever(ss.Settings.from_env())
    got = ids(retriever.search("*", os.environ[env_user], top=50))
    assert ("doc-restricted" in got) is should_see


@pytest.mark.skipif(not LIVE, reason="set RUN_LIVE_TRIMMING_TESTS=1 to run against the sandbox index")
def test_live_graph_outage_returns_nothing():
    real = ss.default_token_provider()

    def broken_graph(scope: str) -> str:
        if scope == ss.GRAPH_SCOPE:
            raise ConnectionError("simulated Graph outage")
        return real(scope)

    retriever = ss.SecureRetriever(ss.Settings.from_env(), token_provider=broken_graph)
    assert retriever.search("*", os.environ["LIVE_ALLOWED_USER_OID"], top=50) == []