"""add vs upsert, metadata on upsert, and delete-by-url. Runs on a copy ('docs_copy'), not on 'docs'."""
import chromadb

from common import HOST, PORT, COLLECTION

client = chromadb.HttpClient(host=HOST, port=PORT)
if "docs_copy" in [c.name for c in client.list_collections()]:
    client.delete_collection("docs_copy")
src = client.get_collection(COLLECTION)
data = src.get(include=["documents", "metadatas", "embeddings"])
cp = client.create_collection("docs_copy", configuration={"hnsw": {"space": "cosine"}}, embedding_function=None)
cp.upsert(ids=data["ids"], documents=data["documents"], metadatas=data["metadatas"], embeddings=data["embeddings"])
print("copy count:", cp.count())

URL = "https://webcrawlerapi.com/docs/api/crawl"
page = cp.get(where={"url": URL}, include=["documents", "metadatas", "embeddings"])
print(f"chunks for {URL.replace('https://webcrawlerapi.com', '')}: {len(page['ids'])}")
one = page["ids"][0]
old_text = page["documents"][0]

print("\n1. add() with an existing id and new text")
cp.add(ids=[one], documents=["CHANGED BY ADD"], embeddings=[page["embeddings"][0]], metadatas=[page["metadatas"][0]])
got = cp.get(ids=[one])
print("   count:", cp.count(), "| text unchanged:", got["documents"][0] == old_text)

print("\n2. upsert() with the same id and new text")
cp.upsert(ids=[one], documents=["CHANGED BY UPSERT"], embeddings=[page["embeddings"][0]], metadatas=[page["metadatas"][0]])
print("   count:", cp.count(), "| text now:", cp.get(ids=[one])["documents"][0])

print("\n3. upsert() with a smaller metadata dict (url only)")
print("   before:", cp.get(ids=[one])["metadatas"][0])
cp.upsert(ids=[one], documents=["x"], embeddings=[page["embeddings"][0]], metadatas=[{"url": URL, "note": "new"}])
print("   after: ", cp.get(ids=[one])["metadatas"][0])

print("\n4. the page shrinks from 4 chunks to 2")
new = list(zip(page["ids"][:2], page["documents"][:2], page["embeddings"][:2], page["metadatas"][:2]))
cp.upsert(ids=[n[0] for n in new], documents=[n[1] for n in new], embeddings=[n[2] for n in new], metadatas=[n[3] for n in new])
print("   upsert only: count", cp.count(), "| chunks for page:", len(cp.get(where={'url': URL})['ids']), "(stale chunks remain)")
cp.delete(where={"url": URL})
print("   delete(where url):    count", cp.count(), "| chunks for page:", len(cp.get(where={'url': URL})['ids']))
cp.upsert(ids=[n[0] for n in new], documents=[n[1] for n in new], embeddings=[n[2] for n in new], metadatas=[n[3] for n in new])
print("   re-upsert 2 chunks:   count", cp.count(), "| chunks for page:", len(cp.get(where={'url': URL})['ids']))

client.delete_collection("docs_copy")
print("\ncopy deleted; docs count:", src.count())