"""Cross-tenant leak tests for 5 retrieval designs (and 2 cache designs).

Every document carries a canary token naming its owner ("CANARY-t07-0012").
A leak = any returned chunk whose canary belongs to another tenant.
Embeddings: hashed bag-of-words (256 dims), deterministic, no model download.
Tenants deliberately share vocabulary, so every query has near neighbours in
other tenants' data (the hard case from the metadata filtering post).
"""
import platform, zlib, numpy as np

T, DOCS_PER_T, DIM, K = 20, 200, 256, 5
rng = np.random.default_rng(0)
VOCAB = [f"w{i}" for i in range(400)]  # shared by all tenants

def embed(text):
    v = np.zeros(DIM, np.float32)
    for tok in text.lower().split():
        if tok.startswith("canary"):
            continue  # canary is payload, not signal
        h = zlib.crc32(tok.encode())
        v[h % DIM] += 1 if (h >> 8) & 1 else -1
    n = np.linalg.norm(v)
    return v / n if n else v

docs = []  # (tenant, text)
for t in range(T):
    for j in range(DOCS_PER_T):
        words = " ".join(rng.choice(VOCAB, 25))
        docs.append((f"t{t:02d}", f"{words} CANARY-t{t:02d}-{j:04d}"))
tenant_of = np.array([d[0] for d in docs])
X = np.stack([embed(d[1]) for d in docs])

def owner(text): return text.split("CANARY-")[1][:3]

# ---- designs: each takes (tenant, query_vec) and returns list of doc indexes
def d_no_filter(tenant, q):
    return list(np.argsort(-(X @ q))[:K])

def d_prompt_only(tenant, q):
    # "You may only use documents of tenant X" lives in the LLM prompt;
    # retrieval itself is unfiltered, so the chunks reach the model.
    return list(np.argsort(-(X @ q))[:K])

def d_post_filter(tenant, q, fetch=50):
    top = np.argsort(-(X @ q))[:fetch]
    return [i for i in top if tenant_of[i] == tenant][:K]

def d_falsy_bug(tenant, q):
    # classic: `if tenant:` skips the filter for "" / None / missing header
    s = X @ q
    if tenant:
        s = np.where(tenant_of == tenant, s, -np.inf)
    top = np.argsort(-s)[:K]
    return [i for i in top if np.isfinite(s[i])]

PART = {t: np.flatnonzero(tenant_of == t) for t in sorted(set(tenant_of))}
def d_partitioned(tenant, q):
    ids = PART.get(tenant)
    if ids is None:
        raise PermissionError("unknown tenant")  # fail closed
    s = X[ids] @ q
    return list(ids[np.argsort(-s)[:K]])

DESIGNS = {"no filter": d_no_filter, "prompt-only rule": d_prompt_only,
           "post-filter top-50": d_post_filter, "filter skipped when tenant falsy": d_falsy_bug,
           "partitioned (pre-filter, fail closed)": d_partitioned}

# ---- test queries: 3 kinds
def queries(tenant):
    others = [i for i in range(len(docs)) if tenant_of[i] != tenant]
    own = np.flatnonzero(tenant_of == tenant)
    r = np.random.default_rng(int(tenant[1:]))
    out = []
    for i in r.choice(own, 30, replace=False):          # a: normal, own content
        out.append(("own-content", embed(docs[i][1])))
    for i in r.choice(others, 30, replace=False):       # b: attacker pastes another tenant's text
        out.append(("copied-other-tenant", embed(docs[i][1])))
    for _ in range(30):                                  # c: generic
        out.append(("generic", embed(" ".join(r.choice(VOCAB, 4)))))
    return out

def run_design(name, fn, tenants):
    res = {k: dict(n=0, leaked_queries=0, leaked_chunks=0, short=0) for k in ("own-content", "copied-other-tenant", "generic")}
    for t in tenants:
        for kind, q in queries(t):
            ids = fn(t, q)
            r = res[kind]; r["n"] += 1
            bad = [i for i in ids if tenant_of[i] != t]
            r["leaked_chunks"] += len(bad); r["leaked_queries"] += bool(bad)
            r["short"] += len(ids) < K
    return res

print(f"python {platform.python_version()} numpy {np.__version__}")
print(f"{T} tenants x {DOCS_PER_T} docs = {len(docs)} chunks, dim {DIM}, k={K}, 90 queries per tenant ({90*T} per design)")
print("\n== retrieval designs: queries with a leak / total, leaked chunks, queries returning < k results ==")
tenants = sorted(PART)
for name, fn in DESIGNS.items():
    res = run_design(name, fn, tenants)
    tot_l = sum(r["leaked_queries"] for r in res.values()); tot_c = sum(r["leaked_chunks"] for r in res.values())
    tot_n = sum(r["n"] for r in res.values()); tot_s = sum(r["short"] for r in res.values())
    print(f"\n{name}: leaked queries {tot_l}/{tot_n}, leaked chunks {tot_c}, short results {tot_s}/{tot_n}")
    for kind, r in res.items():
        print(f"   {kind:20s} leaked {r['leaked_queries']:3d}/{r['n']}  chunks {r['leaked_chunks']:4d}  short {r['short']:3d}")

print("\n== edge-case auth inputs against 'filter skipped when tenant falsy' vs 'partitioned' ==")
q = embed(docs[5][1])
for label, tenant in [("tenant=''", ""), ("tenant=None", None), ("tenant='t99' (unknown)", "t99")]:
    out = {}
    for name in ("filter skipped when tenant falsy", "partitioned (pre-filter, fail closed)"):
        try: ids = DESIGNS[name](tenant, q); out[name] = f"returned {len(ids)} chunks, {len(set(tenant_of[ids]))} tenants"
        except PermissionError: out[name] = "PermissionError"
    print(f"   {label:24s} falsy-bug: {out['filter skipped when tenant falsy']:36s} partitioned: {out['partitioned (pre-filter, fail closed)']}")

print("\n== semantic/answer cache ==")
def sim(a, b): return float(a @ b)
class Cache:
    def __init__(self, per_tenant): self.per_tenant, self.items = per_tenant, []
    def get(self, tenant, q, thr=0.9):
        for t, cq, val in self.items:
            if (not self.per_tenant or t == tenant) and sim(cq, q) >= thr: return val
    def put(self, tenant, q, val): self.items.append((tenant, q, val))
for per_tenant in (False, True):
    c = Cache(per_tenant); hits = leaks = 0
    for t in tenants:  # tenant populates cache with its own retrieval result
        qv = embed(docs[PART[t][0]][1]); c.put(t, qv, [docs[i][1] for i in d_partitioned(t, qv)])
    for t in tenants:  # every tenant then asks the same text as the NEXT tenant's cached query
        nxt = tenants[(tenants.index(t) + 1) % T]
        qv = embed(docs[PART[nxt][0]][1])
        v = c.get(t, qv)
        if v is not None:
            hits += 1; leaks += any(owner(x) != t for x in v)
    print(f"   cache keyed by {'(tenant, query)' if per_tenant else 'query only   '}: hits {hits}/{T}, leaked answers {leaks}/{T}")