"""
测试新的中文 BM25 检索器
验证 jieba 分词是否正常工作
"""
from src.core import RAGSystem

print("=== 测试中文 BM25 检索器 ===\n")

# 创建 RAG 系统实例
rag = RAGSystem()

# 先加载向量索引以获取节点
print("1. 加载向量索引...")
rag.load_index()

if rag.index:
    print("   ✅ 向量索引加载成功\n")
    
    # 从向量索引获取所有节点
    print("2. 从向量索引获取文档节点...")
    from llama_index.core.storage.docstore import SimpleDocumentStore
    
    docstore = rag.index.docstore
    all_nodes = list(docstore.docs.values())
    print(f"   找到 {len(all_nodes)} 个文档节点\n")
    
    # 用中文 BM25 检索器重建索引
    print("3. 构建中文 BM25 检索器...")
    bm25_retriever = rag.bm25_store_manager.setup_bm25(nodes=all_nodes)
    
    if bm25_retriever:
        print("   ✅ 中文 BM25 检索器构建成功\n")
        
        # 测试检索
        print("4. 测试检索:")
        test_query = "什么是RAG"
        print(f"   查询: {test_query}\n")
        
        nodes = bm25_retriever.retrieve(test_query)
        print(f"   检索到 {len(nodes)} 个结果\n")
        
        if nodes:
            print(f"   Top 3 结果:")
            for i, node in enumerate(nodes[:3], 1):
                score = node.score if hasattr(node, 'score') else 'N/A'
                text_preview = node.text[:100] if hasattr(node, 'text') else 'N/A'
                print(f"\n   [{i}] 分数: {score:.4f}")
                print(f"       内容: {text_preview}...")
                
            # 检查是否有非零分数
            non_zero_scores = [n.score for n in nodes if hasattr(n, 'score') and n.score > 0]
            if non_zero_scores:
                print(f"\n   ✅ 成功！检索到 {len(non_zero_scores)} 个非零分数结果")
                print(f"   最高分数: {max(non_zero_scores):.4f}")
            else:
                print(f"\n   ⚠️ 所有结果的分数仍然是 0")
        else:
            print("   ❌ 没有检索到任何结果")
        
        # 测试另一个查询
        print("\n5. 测试另一个查询:")
        test_query2 = "向量数据库"
        print(f"   查询: {test_query2}\n")
        
        nodes2 = bm25_retriever.retrieve(test_query2)
        print(f"   检索到 {len(nodes2)} 个结果")
        
        if nodes2 and len(nodes2) > 0:
            top_node = nodes2[0]
            print(f"   Top 1: 分数 = {top_node.score:.4f}")
            print(f"   内容: {top_node.text[:100]}...")
        
        print("\n" + "="*50)
        print("✅ 中文 BM25 检索器测试完成！")
    else:
        print("   ❌ 中文 BM25 检索器构建失败")
else:
    print("   ❌ 向量索引加载失败")
