"""Create the collection, add payload indexes, upsert every chunk with a payload."""
import uuid
from collections import Counter
from qdrant_client import QdrantClient, models
from common import COLLECTION, read_chunks, section_of

client = QdrantClient(url="http://localhost:6333")
rows, model, dim = read_chunks()
print(f"{len(rows)} chunks, model={model}, dim={dim}")
print("sections:", dict(Counter(section_of(r["url"]) for r in rows).most_common()))

if client.collection_exists(COLLECTION):
    client.delete_collection(COLLECTION)
client.create_collection(
    COLLECTION,
    vectors_config=models.VectorParams(size=dim, distance=models.Distance.COSINE),
)

# Keyword indexes on the fields we filter by
client.create_payload_index(COLLECTION, "section", models.PayloadSchemaType.KEYWORD)
client.create_payload_index(COLLECTION, "url", models.PayloadSchemaType.KEYWORD)

NS = uuid.UUID("6f1c5a3e-0000-4000-8000-000000000000")  # fixed namespace -> stable ids
points = [
    models.PointStruct(
        id=str(uuid.uuid5(NS, r["id"])),
        vector=r["embedding"],
        payload={
            "chunk_id": r["id"],
            "url": r["url"],
            "heading": r["heading"],
            "section": section_of(r["url"]),
            "text": r["text"],
        },
    )
    for r in rows
]
for i in range(0, len(points), 64):
    client.upsert(COLLECTION, points=points[i : i + 64], wait=True)

print("count (exact):", client.count(COLLECTION, exact=True).count)