import time
import jieba  # 导入 jieba 分词库
from Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
from Milvus.Utils.MilvusConnectionPool import *
from Milvus.Config.MulvusConfig import *
from gensim.models import KeyedVectors

# 1. 加载预训练的 Word2Vec 模型
model_path = MS_MODEL_PATH  # 替换为你的 Word2Vec 模型路径
model = KeyedVectors.load_word2vec_format(model_path, binary=False, limit=MS_MODEL_LIMIT)
print(f"模型加载成功，词向量维度: {model.vector_size}")


# 将文本转换为嵌入向量
def text_to_embedding(text):
    words = jieba.lcut(text)  # 使用 jieba 分词
    print(f"文本: {text}, 分词结果: {words}")
    embeddings = [model[word] for word in words if word in model]
    print(f"有效词向量数量: {len(embeddings)}")
    if embeddings:
        avg_embedding = sum(embeddings) / len(embeddings)
        print(f"生成的平均向量: {avg_embedding[:5]}...")  # 打印前 5 维
        return avg_embedding
    else:
        print("未找到有效词，返回零向量")
        return [0.0] * model.vector_size


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

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

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

# 5. 加载集合到内存
collection_manager.load_collection()
print(f"集合 '{collection_name}' 已加载到内存。")

# 6. 输入一句话
input_text = "小学数学中有哪些模型?"

# 7. 将文本转换为嵌入向量
current_embedding = text_to_embedding(input_text)

# 8. 查询与当前对话最相关的历史对话
start_time = time.time()
search_params = {
    "metric_type": "L2",  # 使用 L2 距离度量方式
    "params": {"nprobe": MS_NPROBE}  # 设置 IVF_FLAT 的 nprobe 参数
}
# 哪些文档查询，哪些不查询，我说了算！
# 这样的话，我就可以打多个标签了！
# expr = "document_id in ['MATH_DATA_1', 'MATH_DATA_2']"
expr = "tags['tags'] == 'MATH_DATA_1'"
results = collection_manager.search(
    current_embedding,
    search_params,
    expr=expr,  # 使用in操作符
    limit=10
)

end_time = time.time()

# 9. 输出查询结果
print("最相关的历史对话:")
if results:
    for hits in results:
        for hit in hits:
            try:
                # 查询非向量字段
                record = collection_manager.query_by_id(hit.id)
                print(f"ID: {hit.id}")
                print(f"标签: {record['tags']}")
                print(f"用户问题: {record['user_input']}")
                print(f"时间: {record['timestamp']}")
                print(f"距离: {hit.distance}")
                print("-" * 40)  # 分隔线
            except Exception as e:
                print(f"查询失败: {e}")
else:
    print("未找到相关历史对话，请检查查询参数或数据。")

# 10. 输出查询耗时
print(f"查询耗时: {end_time - start_time:.4f} 秒")

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

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