from pymilvus import FieldSchema, DataType, utility
from WxMini.Milvus.Config.MulvusConfig import *
from WxMini.Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
from WxMini.Milvus.Utils.MilvusConnectionPool import *

# 1. 使用连接池管理 Milvus 连接
milvus_pool = MilvusConnectionPool(host=MS_HOST, port=MS_PORT, max_connections=MS_MAX_CONNECTIONS)

# 2. 从连接池中获取一个连接
connection = milvus_pool.get_connection()

# 3. 初始化集合管理器
collection_name = MS_COLLECTION_NAME
collection_manager = MilvusCollectionManager(collection_name)

# 4. 判断集合是否存在，存在则删除
if utility.has_collection(collection_name):
    print(f"集合 '{collection_name}' 已存在，正在删除...")
    utility.drop_collection(collection_name)
    print(f"集合 '{collection_name}' 已删除。")

# 5. 定义集合的字段和模式
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),  # 主键字段，自动生成 ID
    FieldSchema(name="session_id", dtype=DataType.VARCHAR, max_length=64),  # 会话 ID
    FieldSchema(name="user_input", dtype=DataType.VARCHAR, max_length=2048),  # 用户问题
    FieldSchema(name="model_response", dtype=DataType.VARCHAR, max_length=2048),  # 大模型反馈结果
    FieldSchema(name="timestamp", dtype=DataType.VARCHAR, max_length=32),  # 时间
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=MS_DIMENSION)  # 向量字段，维度为 200
]
schema_description = "Chat records collection with session ID, user input, model response, and timestamp"

# 6. 创建集合
print(f"正在创建集合 '{collection_name}'...")
collection_manager.create_collection(fields, schema_description)
print(f"集合 '{collection_name}' 创建成功。")

# 7. 释放连接
milvus_pool.release_connection(connection)

# 8. 关闭连接池
milvus_pool.close()