"""
Solution: who may see what, for an assistant that searches company documents. Lab 16's
script 01 (and its secure_rag.py): a permission check, four ways to apply it to search,
a canary check for leaks, and the same check run on permissions that went stale.
"""

K = 4  # the assistant answers from the top 4 documents search finds

MODES = ["none", "prompt-only", "post-filter", "pre-filter"]


def principals(user):
    """Everything a user counts as when a document lists who may read it: "public", their own
    "user:<id>", and each of their groups."""
    return {"public", f"user:{user['id']}", *user["groups"]}


def may_read(user, doc_groups):
    """True if the user may read a document readable by `doc_groups` (any one group is
    enough). A disabled account may read nothing."""
    if user.get("disabled"):
        return False
    return bool(principals(user) & set(doc_groups))


def retrieve(ranking, user, mode, acl, k=K):
    """The documents the model will see, best first.

    ranking: every document id, best match first (the search's order for this question).
    acl:     {doc id: [groups that may read it]}, whichever copy of the permissions you use.
    mode:    "none" and "prompt-only" take the top k of everything (the prompt rule only
             labels them); "post-filter" takes the top k, then drops what the user may not
             read; "pre-filter" ranks only what the user may read and takes its top k."""
    if mode == "pre-filter":
        return [d for d in ranking if may_read(user, acl[d])][:k]
    top = ranking[:k]
    if mode == "post-filter":
        return [d for d in top if may_read(user, acl[d])]
    return top


def leaked(answer, canaries, hidden):
    """Ids of the documents in `hidden` (ones this user may not read) whose canaries appear
    in the answer. canaries: {doc id: [strings only that document holds]}. Case-insensitive."""
    text = answer.lower()
    return [d for d in hidden if any(c.lower() in text for c in canaries.get(d, []))]


def shows_forbidden(shown, user_now, acl_now):
    """True if any shown document is one the user may NOT read under today's permissions."""
    return any(not may_read(user_now, acl_now[d]) for d in shown)