from pymilvus import MilvusClient, DataType


client = MilvusClient(
    uri="http://localhost:19530"
)
#for x in client.list_collections():
#    client.drop_collection(x)




# %%

import os
out = os.listdir('/mergedfs/ftp/')

# %%
for cam in out:
    schema = MilvusClient.create_schema(
        auto_id=False,
        nenable_dynamic_field=False,
    )
    schema.add_field(field_name="primary_id",datatype=DataType.INT64, is_primary=True)
    schema.add_field(field_name="filepath", datatype=DataType.VARCHAR, max_length=128)
    schema.add_field(field_name="frame_number", datatype=DataType.INT32)
    schema.add_field(field_name="date", datatype=DataType.VARCHAR, max_length=len('20241220'), is_partition_key=True)
    schema.add_field(field_name="so400m", datatype=DataType.FLOAT16_VECTOR, dim=1024)


    index_params = client.prepare_index_params()

    index_params.add_index(
        field_name="primary_id",
        index_type="STL_SORT")

    index_params.add_index(
        field_name="filepath",
        index_type="Trie")

    index_params.add_index(
        field_name="so400m", 
        index_type="IVF_SQ8",
        metric_type="COSINE",
        params={'nlist':128})

    index_params.add_index(
        field_name='date',
        index_type='Trie')

    client.create_collection(
        collection_name=f"nuggets_{cam}_so400m_siglip2",
        schema=schema,
        index_params=index_params
    )
    print(cam)

#res = client.get_load_state( collection_name="nuggets_ptz_so400m_siglip2")
#res = client.load_collection(collection_name="nuggets_so400m")