import time
from collections import Counter

import chromadb

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

rows, model, dim = read_chunks()
print(f"{len(rows)} chunks, model={model}, dim={dim}")
print("chunks with an empty heading:", sum(1 for r in rows if not r["heading"]))

client = chromadb.HttpClient(host=HOST, port=PORT)
print("max batch size:", client.get_max_batch_size())
# start clean so the timing is a first load, not a re-run
if COLLECTION in [c.name for c in client.list_collections()]:
    client.delete_collection(COLLECTION)
col = client.create_collection(
    COLLECTION,
    configuration={"hnsw": {"space": "cosine"}},
    embedding_function=None,  # we always pass our own vectors
)
print("embedding function on collection:", col.configuration.get("embedding_function"))

batch = min(client.get_max_batch_size(), 128)
t0 = time.time()
for i in range(0, len(rows), batch):
    part = rows[i : i + batch]
    col.upsert(
        ids=[r["id"] for r in part],
        documents=[r["text"] for r in part],
        embeddings=[r["embedding"] for r in part],
        metadatas=[
            {"url": r["url"], "heading": r["heading"], "section": section_of(r["url"])}
            for r in part
        ],
    )
print(f"upserted in {time.time() - t0:.2f}s, batch size {batch}")

print("count:", col.count())
assert col.count() == len(rows)

got = col.get(include=["metadatas"])
by_section = Counter(m["section"] for m in got["metadatas"])
print(f"{len(by_section)} sections, total {sum(by_section.values())}")
for s, n in by_section.most_common():
    print(f"  {s:20s} {n}")