"""Per-customer isolation with a partition key. Three fake customers share one collection.
Runs on whatever KB_URI points at (Lite by default, also used for standalone and Zilliz)."""
import hashlib
from collections import Counter

from pymilvus import DataType

from common import connect, drop_if_exists, read_chunks, to_record, embed_query

NAME = "docs_tenants"
rows, model, dim = read_chunks()
client = connect()
drop_if_exists(client, NAME)


def tenant(url):  # a page belongs to one customer, chosen by a stable hash of its URL
    return "customer_" + "abc"[int(hashlib.md5(url.encode()).hexdigest(), 16) % 3]


schema = client.create_schema(auto_id=False, enable_dynamic_field=True)
schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=128)
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=dim)
schema.add_field("url", DataType.VARCHAR, max_length=512)
schema.add_field("text", DataType.VARCHAR, max_length=16384)
schema.add_field("customer_id", DataType.VARCHAR, max_length=64, is_partition_key=True)
ip = client.prepare_index_params()
ip.add_index(field_name="vector", index_type="AUTOINDEX", metric_type="COSINE")
try:
    client.create_collection(NAME, schema=schema, index_params=ip, num_partitions=8)
    print("create_collection with is_partition_key=True, num_partitions=8: accepted")
except Exception as e:
    print(f"REJECTED: {type(e).__name__}: {str(e)[:400]}")
    raise SystemExit(0)

records = []
for r in rows:
    rec = to_record(r)
    records.append({"id": rec["id"], "vector": rec["vector"], "url": rec["url"], "text": rec["text"],
                    "customer_id": tenant(rec["url"])})
client.insert(NAME, data=records)
client.load_collection(NAME)
per = Counter(r["customer_id"] for r in records)
print("rows per customer:", dict(sorted(per.items())))
print("list_partitions:", client.list_partitions(NAME))

q = "How do I limit how many pages a crawl fetches?"
qv = embed_query(q, model)
allhits = client.search(NAME, data=[qv], limit=5, output_fields=["url", "customer_id"], consistency_level="Strong")[0]
print(f"\nno filter, top 5 customers: {[h['entity']['customer_id'] for h in allhits]}")
for cust in sorted(per):
    res = client.search(NAME, data=[qv], limit=5, filter=f"customer_id == '{cust}'",
                        output_fields=["url", "customer_id"], consistency_level="Strong")[0]
    ok = len(res) == 5 and all(h["entity"]["customer_id"] == cust for h in res)
    print(f"\n{cust}: {len(res)} results, all belong to {cust}: {ok}")
    for h in res[:2]:
        print(f"  {h['distance']:.3f}  {h['entity']['url'].replace('https://webcrawlerapi.com', '')}")

# what if you forget the filter? the key does not enforce anything on its own
print("\nforgot the filter -> hits from other customers are returned (see the 'no filter' line above)")
drop_if_exists(client, NAME)
client.close()