#!/usr/bin/env python3
"""
测试混合模式和BM25检索功能
"""
import sys
import os
sys.path.append('.')

from src.core.rag_system import RAGSystem

def test_all_retrieval_modes():
    """测试所有检索模式"""
    print("=" * 60)
    print("测试所有检索模式")
    print("=" * 60)
    
    try:
        # 初始化RAG系统
        print("初始化RAG系统...")
        rag = RAGSystem()
        
        # 测试查询
        query = "ai"
        
        # 1. 测试向量检索
        print(f"\n{'='*40}")
        print("1. 测试向量检索 (vector)")
        print(f"{'='*40}")
        result_vector = rag.retrieve(query, mode="vector")
        if result_vector and result_vector.get('documents'):
            print(f"✅ 向量检索成功")
            print(f"   检索到 {len(result_vector['documents'])} 个文档片段")
            print(f"   前3个结果:")
            for i, doc in enumerate(result_vector['documents'][:3]):
                print(f"     {i+1}. 分数: {doc.get('score', 0):.4f}, 内容: {doc.get('text', '')[:80]}...")
        else:
            print("❌ 向量检索失败")
        
        # 2. 测试BM25检索
        print(f"\n{'='*40}")
        print("2. 测试BM25检索 (bm25)")
        print(f"{'='*40}")
        result_bm25 = rag.retrieve(query, mode="bm25")
        if result_bm25 and result_bm25.get('documents'):
            print(f"✅ BM25检索成功")
            print(f"   检索到 {len(result_bm25['documents'])} 个文档片段")
            print(f"   前3个结果:")
            for i, doc in enumerate(result_bm25['documents'][:3]):
                print(f"     {i+1}. 分数: {doc.get('score', 0):.4f}, 内容: {doc.get('text', '')[:80]}...")
        else:
            print("❌ BM25检索失败")
        
        # 3. 测试混合检索
        print(f"\n{'='*40}")
        print("3. 测试混合检索 (hybrid)")
        print(f"{'='*40}")
        result_hybrid = rag.retrieve(query, mode="hybrid")
        if result_hybrid and result_hybrid.get('documents'):
            print(f"✅ 混合检索成功")
            print(f"   检索到 {len(result_hybrid['documents'])} 个文档片段")
            print(f"   检索模式: {result_hybrid.get('mode', 'unknown')}")
            print(f"   前3个结果:")
            for i, doc in enumerate(result_hybrid['documents'][:3]):
                source = doc.get('source', 'unknown')
                print(f"     {i+1}. 分数: {doc.get('score', 0):.4f}, 来源: {source}, 内容: {doc.get('text', '')[:80]}...")
        else:
            print("❌ 混合检索失败")
        
        # 4. 比较不同检索模式的结果
        print(f"\n{'='*40}")
        print("4. 检索模式比较")
        print(f"{'='*40}")
        
        if result_vector and result_bm25 and result_hybrid:
            print(f"向量检索结果数: {len(result_vector['documents'])}")
            print(f"BM25检索结果数: {len(result_bm25['documents'])}")
            print(f"混合检索结果数: {len(result_hybrid['documents'])}")
            
            # 检查混合检索是否包含两种来源
            sources = set()
            if result_hybrid['documents']:
                for doc in result_hybrid['documents']:
                    sources.add(doc.get('source', 'unknown'))
            
            print(f"混合检索来源: {sorted(sources)}")
            
            if 'vector' in sources and 'bm25' in sources:
                print("✅ 混合检索成功融合了向量和BM25结果")
            elif len(sources) > 0:
                print(f"⚠️ 混合检索只包含 {sources} 来源")
            else:
                print("❌ 混合检索没有来源信息")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_cli_retrieval():
    """测试CLI检索命令"""
    print(f"\n{'='*60}")
    print("测试CLI检索命令")
    print(f"{'='*60}")
    
    commands = [
        ("向量检索", "python cli.py retrieve -q ai --mode vector"),
        ("BM25检索", "python cli.py retrieve -q ai --mode bm25"),
        ("混合检索", "python cli.py retrieve -q ai --mode hybrid"),
    ]
    
    for name, cmd in commands:
        print(f"\n{'='*40}")
        print(f"测试: {name}")
        print(f"命令: {cmd}")
        print(f"{'='*40}")
        
        # 执行命令
        os.system(cmd)
        
        print(f"\n{name} 测试完成")

def test_query_with_llm():
    """测试使用LLM的问答功能"""
    print(f"\n{'='*60}")
    print("测试使用LLM的问答功能")
    print(f"{'='*60}")
    
    try:
        # 初始化RAG系统
        print("初始化RAG系统...")
        rag = RAGSystem()
        
        # 测试查询
        query = "什么是人工智能？"
        print(f"\n查询: {query}")
        
        # 使用LLM生成答案
        print("使用LLM生成答案...")
        result = rag.query(query)
        
        if result:
            print(f"✅ 问答成功")
            print(f"答案: {result.get('answer', '')[:200]}...")
            print(f"检索到的文档片段数: {len(result.get('source_documents', []))}")
        else:
            print("❌ 问答失败")
            
    except Exception as e:
        print(f"❌ 问答测试失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    print("开始测试混合模式和BM25检索功能...")
    
    # 更新TODO状态
    print("\n[7/9] 修复BM25检索器加载问题 - ✅ 已完成")
    
    # 测试所有检索模式
    print("\n[8/9] 测试混合模式和BM25检索功能...")
    if test_all_retrieval_modes():
        print("\n✅ 混合模式和BM25检索功能测试通过")
    
    # 测试CLI命令
    test_cli_retrieval()
    
    # 测试问答功能
    test_query_with_llm()
    
    print(f"\n{'='*60}")
    print("测试完成")
    print(f"{'='*60}")