"""Self-contained quickstart: load chunks, create the vector index, run one $vectorSearch. Uses its own database, dropped at the end."""
import json, time
from pymongo import MongoClient
from pymongo.operations import SearchIndexModel

client = MongoClient("mongodb://mongo:27017/?directConnection=true")  # from the host: mongodb://127.0.0.1:27017/?directConnection=true
col = client["quickstart"]["chunks"]
col.drop()
rows = [json.loads(line) for line in open("/data/chunks.jsonl")]           # id, url, heading, text, embedding (768 floats)
col.insert_many([{"_id": r["id"], "url": r["url"], "heading": r["heading"], "text": r["text"], "embedding": r["embedding"]} for r in rows])

definition = {"fields": [{"type": "vector", "path": "embedding", "numDimensions": 768, "similarity": "cosine"}]}
col.create_search_index(SearchIndexModel(definition=definition, name="qs_vec", type="vectorSearch"))
while not any(i["name"] == "qs_vec" and i.get("queryable") for i in col.list_search_indexes()):
    time.sleep(0.25)                                                         # PENDING -> READY

question = "How do I limit how many pages a crawl fetches?"
vec = json.load(open("qvecs.json"))[question]                                # embedded earlier with nomic-embed-text
hits = col.aggregate([
    {"$vectorSearch": {"index": "qs_vec", "path": "embedding", "queryVector": vec, "numCandidates": 100, "limit": 3}},
    {"$project": {"url": 1, "heading": 1, "score": {"$meta": "vectorSearchScore"}}},
])
for h in hits:
    print(f"{h['score']:.4f}  {h['url']}  | {h['heading']}")
client.drop_database("quickstart")