"""Create the DocsChunk collection and load chunks.jsonl with our own vectors."""
import json
import weaviate
from weaviate.classes.config import Configure, DataType, Property, Tokenization, VectorDistances

rows = [json.loads(line) for line in open("../data/chunks.jsonl")]

with weaviate.connect_to_local(port=8080, grpc_port=50051) as client:
    if client.collections.exists("DocsChunk"):
        client.collections.delete("DocsChunk")  # only our own collection, safe to re-run
    docs = client.collections.create(
        "DocsChunk",
        vectorizer_config=Configure.Vectorizer.none(),
        vector_index_config=Configure.VectorIndex.hnsw(distance_metric=VectorDistances.COSINE),
        properties=[
            Property(name="chunk_id", data_type=DataType.TEXT, tokenization=Tokenization.FIELD),
            Property(name="url", data_type=DataType.TEXT, tokenization=Tokenization.FIELD),
            Property(name="heading", data_type=DataType.TEXT),
            Property(name="text", data_type=DataType.TEXT),
        ],
    )
    with docs.batch.dynamic() as batch:
        for r in rows:
            batch.add_object(
                properties={"chunk_id": r["id"], "url": r["url"], "heading": r["heading"], "text": r["text"]},
                vector=r["embedding"],
            )
    print("sent:", len(rows))
    print("failed objects:", len(docs.batch.failed_objects))
    print("count in Weaviate:", docs.aggregate.over_all(total_count=True).total_count)