import time
from collections import Counter

from pymilvus import DataType

from common import connect, drop_if_exists, COLLECTION, read_chunks, to_record

rows, model, dim = read_chunks()
print(f"{len(rows)} chunks, model={model}, dim={dim}")
records = [to_record(r) for r in rows]
longest = max(len(r["text"]) for r in records)
print(f"longest chunk text: {longest} characters")

client = connect()
drop_if_exists(client, COLLECTION)

# Explicit schema. "model" is NOT declared: it goes into the dynamic field.
schema = client.create_schema(auto_id=False, enable_dynamic_field=True)
schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=128)
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=dim)
schema.add_field("url", DataType.VARCHAR, max_length=512)
schema.add_field("heading", DataType.VARCHAR, max_length=512)
schema.add_field("section", DataType.VARCHAR, max_length=64)
schema.add_field("text", DataType.VARCHAR, max_length=16384)

index_params = client.prepare_index_params()
index_params.add_index(field_name="vector", index_type="AUTOINDEX", metric_type="COSINE")

t0 = time.time()
client.create_collection(COLLECTION, schema=schema, index_params=index_params)
res = client.insert(COLLECTION, data=records)
print(f"created collection and inserted in {time.time() - t0:.2f}s; insert_count: {res['insert_count']}")

stats = client.get_collection_stats(COLLECTION)
print("stats:", stats)
# Strong consistency: on a server or Zilliz the default (Bounded) count can still read 0 right after an insert
count = client.query(COLLECTION, filter="", output_fields=["count(*)"], consistency_level="Strong")
print("count(*):", count[0])
assert count[0]["count(*)"] == len(rows)

print("\ncollection description:")
d = client.describe_collection(COLLECTION)
for f in d["fields"]:
    print(f"  {f['name']:8s} {DataType(f['type']).name:14s} {f.get('params', {})}{' PK' if f.get('is_primary') else ''}")
print("  enable_dynamic_field:", d.get("enable_dynamic_field"))
print("index:", client.describe_index(COLLECTION, "vector"))

sections = Counter(r["section"] for r in client.query(COLLECTION, filter="", output_fields=["section"], limit=1000))
print(f"\n{len(sections)} sections, total {sum(sections.values())}")
for s, n in sections.most_common():
    print(f"  {s:20s} {n}")

# dynamic field: filter on a key that is not in the schema
dyn = client.query(COLLECTION, filter='model == "nomic-embed-text"', output_fields=["id", "model"], limit=3)
print("\ndynamic field filter model == \"nomic-embed-text\", first 3:", [dict(d) for d in dyn])
n = client.query(COLLECTION, filter='model == "nomic-embed-text"', output_fields=["count(*)"])
print("rows matching:", n[0]["count(*)"])
client.close()