import os
import logging
from logging.handlers import RotatingFileHandler

import jieba
from gensim.models import KeyedVectors

from Config.Config import ES_CONFIG, MS_MODEL_PATH, MS_MODEL_LIMIT
from ElasticSearch.Utils.ElasticsearchConnectionPool import ElasticsearchConnectionPool

# 初始化日志
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# 确保日志目录存在
os.makedirs('Logs', exist_ok=True)
handler = RotatingFileHandler('Logs/start.log', maxBytes=1024 * 1024, backupCount=5)
handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)

# 1. 加载预训练的 Word2Vec 模型
model = KeyedVectors.load_word2vec_format(MS_MODEL_PATH, binary=False, limit=MS_MODEL_LIMIT)
logger.info(f"模型加载成功，词向量维度: {model.vector_size}")



def init_es_pool():
    # 初始化Elasticsearch连接池
    es_pool = ElasticsearchConnectionPool(
        hosts=ES_CONFIG["hosts"],
        basic_auth=ES_CONFIG["basic_auth"],
        verify_certs=ES_CONFIG["verify_certs"],
        max_connections=50
    )
    logger.info("Elasticsearch连接池初始化完成")
    return es_pool


# 导入EsSearchUtil
from Util.EsSearchUtil import EsSearchUtil

# 初始化EsSearchUtil
es_search_util = EsSearchUtil(ES_CONFIG)


def main():
    # 初始化ES连接池
    es_pool = init_es_pool()
    
    # 测试查询
    query = "小学数学中有哪些模型"
    query_tags = ["MATH_1"]  # 默认搜索标签，可修改
    print(f"\n=== 开始执行查询 ===")
    print(f"原始查询文本: {query}")

    # 执行混合搜索
    es_conn = es_pool.get_connection()
    try:
        # 向量搜索
        print("\n=== 向量搜索阶段 ===")
        print("1. 文本分词和向量化处理中...")
        query_embedding = es_search_util.text_to_embedding(query)
        print(f"2. 生成的查询向量维度: {len(query_embedding)}")
        print(f"3. 前5维向量值: {query_embedding[:5]}")
        
        print("4. 正在执行Elasticsearch向量搜索...")
        vector_results = es_conn.search(
            index=ES_CONFIG['index_name'],
            body={
                "query": {
                    "script_score": {
                        "query": {
                            "bool": {
                                "should": [
                                    {
                                        "terms": {
                                            "tags.tags": query_tags
                                        }
                                    }
                                ],
                                "minimum_should_match": 1
                            }
                        },
                        "script": {
                            "source": "double score = cosineSimilarity(params.query_vector, 'embedding'); return score >= 0 ? score : 0",
                            "params": {"query_vector": query_embedding}
                        }
                    }
                },
                "size": 5
            }
        )
        print(f"5. 向量搜索结果数量: {len(vector_results['hits']['hits'])}")
        
        # 文本精确搜索
        print("\n=== 文本精确搜索阶段 ===")
        print("1. 正在执行Elasticsearch文本精确搜索...")
        text_results = es_conn.search(
            index=ES_CONFIG['index_name'],
            body={
                "query": {
                    "bool": {
                        "must": [
                            {
                                "match": {
                                    "user_input": query
                                }
                            },
                            {
                                "terms": {
                                    "tags.tags": query_tags
                                }
                            }
                        ]
                    }
                },
                "size": 5
            }
        )
        print(f"2. 文本搜索结果数量: {len(text_results['hits']['hits'])}")
        
        # 打印详细结果
        print("\n=== 最终搜索结果 ===")
        print(f" 向量搜索结果: {len(vector_results['hits']['hits'])}条")
        for i, hit in enumerate(vector_results['hits']['hits'], 1):
            print(f"   {i}. 文档ID: {hit['_id']}, 相似度分数: {hit['_score']:.2f}")
            print(f"      内容: {hit['_source']['user_input']}")
            # print(f"      详细: {hit['_source']['tags']['full_content']}")
            
        print("\n文本精确搜索结果:")
        for i, hit in enumerate(text_results['hits']['hits']):
            print(f"  {i+1}. 文档ID: {hit['_id']}, 匹配分数: {hit['_score']:.2f}")
            print(f"      内容: {hit['_source']['user_input']}")
            # print(f"      详细: {hit['_source']['tags']['full_content']}")
        
    finally:
        es_pool.release_connection(es_conn)
    
    # 关闭连接池
    es_pool.close()

if __name__ == "__main__":
    main()