import traceback

from pymilvus import connections,utility,FieldSchema,CollectionSchema,DataType,Collection,Milvus

current_host = None
def init_milvus(host):
    global current_host
    current_host = host
    connections.connect("default",host=host,port=19530)

def drop_embedding_collection(collection_name):
    print("drop collection:",collection_name)
    utility.drop_collection(collection_name)


def create_embedding_schema(collection_name,fields,index_name,index_params):
    has = utility.has_collection(collection_name)
    if not has:
        print("creating collection")
        coll_schema = CollectionSchema(fields,"this is the embedding schema")
        coll = Collection(collection_name,coll_schema)

        #create index for milvus_embedding
        coll.create_index(index_name,index_params=index_params)

def getCollection(collection_name):
    try:
        coll = Collection(collection_name)
        print(collection_name,"num_entities",coll.num_entities)
        return coll
    except Exception as e:
        init_milvus(current_host)
        coll = Collection(collection_name)
        print(collection_name,"num_entities",coll.num_entities)
        return coll


def insert_embedding(coll,entities,retry_times =3):
    for _ in range(retry_times):
        try:
            coll.insert(entities)
            coll.load()
            print("num_entities",coll.num_entities)
            return
        except Exception as e:
            init_milvus(current_host)


def search_embedding(coll,index_name,vector,search_params,output_fields,limit=3,retry_times=3):
    for _ in range(retry_times):
        try:
            list_result = []
            result = coll.search(vector,index_name,search_params,top_k=limit,output_fields=output_fields,limit=limit)
            for hits in result:
                for hit in hits:
                    list_result.append(hit)

            return list_result
        except Exception as e:
            init_milvus(current_host)

if __name__ == '__main__':
    # drop_embedding_collection()
    create_embedding_schema()

