from pymilvus import CollectionSchema, FieldSchema, DataType,client
from pymilvus import Collection, db, connections,utility,Function,FunctionType
from pymilvus import (
    MilvusClient, DataType
)

def new_create_coll():
    client = MilvusClient(
        uri="http://172.18.107.78:19530",
        token="root:Milvus"
    )
    client.use_database("JXTest")

    schema = MilvusClient.create_schema(
        auto_id=True,
        enable_dynamic_field=True,
    )

    analyzer_params_custom = {
    "tokenizer": "jieba",
    "type": "chinese",
    "filter":  ["cnalphanumonly"]
}

    schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
    schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1000,enable_analyzer=True,analyzer_params=analyzer_params_custom)
    schema.add_field(field_name="sparse", datatype=DataType.SPARSE_FLOAT_VECTOR)
    schema.add_field(field_name="dense", datatype=DataType.FLOAT_VECTOR, dim=5)

    bm25_function = Function(
        name="text_bm25_emb", # Function name
        input_field_names=["text"], # Name of the VARCHAR field containing raw text data
        output_field_names=["sparse"], # Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
        function_type=FunctionType.BM25,
    )

    schema.add_function(bm25_function)

    index_params = client.prepare_index_params()

    # Add indexes
    index_params.add_index(
        field_name="dense",
        index_name="dense_index",
        index_type="AUTOINDEX",#"IVF_FLAT"
        metric_type="COSINE", #L2
        params={"nlist": 128},
    )

    index_params.add_index(
        field_name="sparse",
        index_name="sparse_index",
        index_type="SPARSE_INVERTED_INDEX",  # Index type for sparse vectors
        metric_type="BM25",  # Set to `BM25` when using function to generate sparse vectors
        params={"inverted_index_algo": "DAAT_MAXSCORE"},  # The ratio of small vector values to be dropped during indexing
    )

    client.create_collection(
    collection_name="hybrid_search_collection",
    schema=schema,
    index_params=index_params
    )
    return "Sucessful"



def create_collections(collections_name,database_name : str ="JXTest"):
  db.using_database(database_name)

  id = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True)
  text = FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=4096,enable_analyzer=True)
  header_1 = FieldSchema(name="header_1", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
  header_2 = FieldSchema(name="header_2", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
  header_3 = FieldSchema(name="header_3", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
  vector = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024,)
  sparse = FieldSchema(name="sparse_bm25", dtype=DataType.SPARSE_FLOAT_VECTOR)
  schema = CollectionSchema(
      fields=[id, header_1, header_2, header_3, text, vector,sparse],
      auto_id=True,
      enable_dynamic_field=True
  )
  index_params = {
    "field_name": "vector",
    "index_name": "vector_index",
    "metric_type": "L2",
    "index_type": "IVF_FLAT", #FLAT
    "params": {"nlist": 128},
    "field_name":"sparse_bm25",  
    "index_name":"sparse_bm25_index",  
    "index_type":"SPARSE_WAND",   
    "metric_type":"BM25",
    "params":{
        "inverted_index_algo": "DAAT_MAXSCORE", # Algorithm for building and querying the index. Valid values: DAAT_MAXSCORE, DAAT_WAND, TAAT_NAIVE.
        "bm25_k1": 1.2,
        "bm25_b": 0.75
    }

  }
  bm25_function = Function(  
        name="bm25",  
        function_type=FunctionType.BM25,  
        input_field_names=["text"],  
        output_field_names="sparse_bm25",  
    )  
  schema.add_function(bm25_function)  

  collection = Collection(name=collections_name, schema=schema, using='default', shards_num=2)
  collection.create_index(field_name="vector", index_params=index_params)

  utility.index_building_progress("Hybrid")
  return "collections created"


if __name__ == "__main__":
    # conn = connections.connect(host="172.18.107.78",port=19530,user="user",password="Milvus")
    # res = create_collections("HybirdTest","JXTest")
    # print(res)
    res = new_create_coll()