"""
测试BM25检索功能修复
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.config import get_config_manager
from src.storage import BM25StoreManager
from src.tokenizers import create_tokenizer

def test_bm25_retrieval():
    """测试BM25检索功能"""
    print("=== 测试 BM25 检索功能 ===\n")
    
    # 获取配置
    config = get_config_manager()
    
    print(f"1. 配置检查:")
    print(f"   检索模式: {config.retrieval.mode}")
    print(f"   BM25启用: {config.retrieval.use_bm25}")
    print(f"   BM25语言: {config.retrieval.bm25_language}")
    print(f"   BM25存储目录: {config.paths.bm25_storage_dir}")
    
    # 创建BM25存储管理器
    bm25_manager = BM25StoreManager(config=config)
    
    print(f"\n2. 检查BM25检索器状态:")
    print(f"   BM25检索器已加载: {bm25_manager.is_loaded()}")
    
    # 尝试获取BM25检索器
    bm25_retriever = bm25_manager.get_bm25_retriever()
    if bm25_retriever:
        print(f"   ✅ BM25检索器获取成功")
        
        # 测试检索
        print(f"\n3. 测试检索:")
        test_queries = ["什么是RAG", "向量数据库", "人工智能"]
        
        for query in test_queries:
            print(f"\n   查询: '{query}'")
            try:
                nodes = bm25_retriever.retrieve(query)
                print(f"   检索到 {len(nodes)} 个结果")
                
                if nodes:
                    for i, node in enumerate(nodes[:2], 1):
                        score = node.score if hasattr(node, 'score') else 'N/A'
                        text = node.node.get_content() if hasattr(node.node, 'get_content') else str(node.node)[:100]
                        print(f"      [{i}] 分数: {score}")
                        print(f"          内容: {text[:80]}...")
                else:
                    print("      ⚠️ 未检索到结果")
            except Exception as e:
                print(f"      ❌ 检索失败: {e}")
    else:
        print(f"   ❌ BM25检索器未加载")
        
        # 尝试手动加载
        print(f"\n4. 尝试手动加载BM25检索器...")
        try:
            # 创建分词器
            tokenizer = create_tokenizer(
                language=config.retrieval.bm25_language,
                use_stemmer=(config.retrieval.bm25_language == "english")
            )
            
            # 尝试从持久化目录加载
            from src.retrievers import ChineseBM25Retriever
            bm25_retriever = ChineseBM25Retriever.from_persist_dir(
                persist_dir=str(config.paths.bm25_storage_dir),
                tokenizer=tokenizer,
                verbose=True
            )
            
            if bm25_retriever:
                print(f"   ✅ 手动加载成功")
                bm25_manager.bm25_retriever = bm25_retriever
                
                # 测试检索
                print(f"\n5. 测试检索:")
                test_query = "什么是RAG"
                print(f"   查询: '{test_query}'")
                
                nodes = bm25_retriever.retrieve(test_query)
                print(f"   检索到 {len(nodes)} 个结果")
                
                if nodes:
                    for i, node in enumerate(nodes[:2], 1):
                        score = node.score if hasattr(node, 'score') else 'N/A'
                        text = node.node.get_content() if hasattr(node.node, 'get_content') else str(node.node)[:100]
                        print(f"      [{i}] 分数: {score}")
                        print(f"          内容: {text[:80]}...")
                else:
                    print("      ⚠️ 未检索到结果")
            else:
                print(f"   ❌ 手动加载失败")
                
        except Exception as e:
            print(f"   ❌ 加载失败: {e}")
            import traceback
            traceback.print_exc()
    
    print(f"\n" + "="*50)
    print(f"测试完成")

if __name__ == "__main__":
    test_bm25_retrieval()