from pymilvus import MilvusClient, DataType
import os
from tqdm import tqdm
client = MilvusClient(
    uri="http://localhost:19530"
)


paths_scan = ['/srv/ftp_tcc','/srv/ftp','/srv/ftp_local/']

outff = list()
for xpp in paths_scan:
    outff.extend([x.path for x in os.scandir(xpp) if x.is_dir()])




out = [x.split('/')[3] for x in outff]

# %%
c_collections = set(client.list_collections())
for cam in tqdm(out):
    c_col_name = f"nuggets_{cam}_so400m_siglip2"
    if c_col_name in c_collections:
        print(f"Skipping {c_col_name}")
        continue
    else:
        print(f"Creating {c_col_name}")
    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=1152)


    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="DISKANN",
        metric_type="COSINE",
        params={'nlist':128})

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

    client.create_collection(
        collection_name=c_col_name,
        schema=schema,
        index_params=index_params
    )

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

#import numpy

#c = numpy.load('/srv/ftp_tcc/leopards1/2025/07/18/Leopards1_00_20250718170024.oclip_embeds.npz')