import chromadb

from common import HOST, PORT, COLLECTION, read_chunks, embed_query

rows, model, dim = read_chunks()
col = chromadb.HttpClient(host=HOST, port=PORT).get_collection(COLLECTION)
Q = "How do I get notified when a crawl job finishes?"
vec = embed_query(Q, model)


def show(label, **kw):
    res = col.query(query_embeddings=[vec], n_results=5, include=["metadatas", "distances"], **kw)
    print(f"\n{label}")
    for m, d in zip(res["metadatas"][0], res["distances"][0]):
        print(f"  {d:.3f}  {m['section']:15s} {m['url'].replace('https://webcrawlerapi.com', '')}")
    return res["metadatas"][0]


show("no filter")
r = show("where section == sdk", where={"section": "sdk"})
assert all(m["section"] == "sdk" for m in r)
r = show("where section == guides", where={"section": "guides"})
assert all(m["section"] == "guides" for m in r)
r = show("where section $in [guides, errors]", where={"section": {"$in": ["guides", "errors"]}})
assert all(m["section"] in ("guides", "errors") for m in r)

print("\nhits per filter (get, no query):")
WITH_CRAWL = {"$and": [{"section": "api"}, {"url": {"$ne": "https://webcrawlerapi.com/docs/api/crawl"}}]}
filters = [
    ("section == api", {"section": "api"}),
    ("section == sdk", {"section": "sdk"}),
    ("section $in [guides, errors]", {"section": {"$in": ["guides", "errors"]}}),
    ("$or section api / sdk", {"$or": [{"section": "api"}, {"section": "sdk"}]}),
    ("section $ne api", {"section": {"$ne": "api"}}),
    ("section api AND url != /docs/api/crawl", WITH_CRAWL),
]
for label, w in filters:
    print(f"  {len(col.get(where=w, include=[])['ids']):4d}  {label}")

print("\nwhere_document:")
for needle in ["max_depth", "Max_Depth", "webhook"]:
    n = len(col.get(where_document={"$contains": needle}, include=[])["ids"])
    print(f"  $contains {needle!r}: {n} chunks")
vec2 = embed_query("How do I limit how many pages a crawl fetches?", model)
res = col.query(query_embeddings=[vec2], n_results=5, where_document={"$contains": "max_depth"},
                include=["metadatas", "distances", "documents"])
print("  query + where_document $contains 'max_depth':")
for m, d, t in zip(res["metadatas"][0], res["distances"][0], res["documents"][0]):
    assert "max_depth" in t
    print(f"    {d:.3f}  {m['url'].replace('https://webcrawlerapi.com', '')}")
res = col.query(query_embeddings=[vec2], n_results=5, where={"section": "sdk"}, where_document={"$contains": "max_depth"},
                include=["metadatas", "distances"])
print("  same + where section == sdk:", [m["url"].replace("https://webcrawlerapi.com", "") for m in res["metadatas"][0]])

print("\nget(where=section errors, limit=3):")
g = col.get(where={"section": "errors"}, limit=3, include=["metadatas"])
for m in g["metadatas"]:
    print("  ", m["heading"])
page2 = col.get(where={"section": "errors"}, limit=3, offset=3, include=["metadatas"])
print("   offset=3 ->", len(page2["ids"]), "more")