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

from src.core.rag_system import RAGSystem
from src.config import Config

def test_bm25_retrieval():
    """测试BM25检索功能"""
    print("=" * 60)
    print("测试BM25检索功能")
    print("=" * 60)
    
    try:
        # 初始化RAG系统
        print("初始化RAG系统...")
        rag = RAGSystem()
        
        # 测试BM25检索
        print("\n测试BM25检索...")
        result = rag.retrieve("ai", mode="bm25")
        
        if result:
            print(f"✅ BM25检索成功")
            print(f"检索到 {len(result.get('documents', []))} 个文档片段")
            
            # 显示前3个结果
            for i, doc in enumerate(result.get('documents', [])[:3]):
                print(f"\n结果 {i+1}:")
                print(f"  分数: {doc.get('score', 0):.4f}")
                print(f"  内容: {doc.get('text', '')[:100]}...")
        else:
            print("❌ BM25检索失败")
            
    except Exception as e:
        print(f"❌ BM25检索测试失败: {e}")
        import traceback
        traceback.print_exc()

def test_hybrid_retrieval():
    """测试混合检索功能"""
    print("\n" + "=" * 60)
    print("测试混合检索功能")
    print("=" * 60)
    
    try:
        # 初始化RAG系统
        print("初始化RAG系统...")
        rag = RAGSystem()
        
        # 测试混合检索
        print("\n测试混合检索...")
        result = rag.retrieve("ai", mode="hybrid")
        
        if result:
            print(f"✅ 混合检索成功")
            print(f"检索到 {len(result.get('documents', []))} 个文档片段")
            print(f"检索模式: {result.get('mode', 'unknown')}")
            
            # 显示前3个结果
            for i, doc in enumerate(result.get('documents', [])[:3]):
                print(f"\n结果 {i+1}:")
                print(f"  分数: {doc.get('score', 0):.4f}")
                print(f"  来源: {doc.get('source', 'unknown')}")
                print(f"  内容: {doc.get('text', '')[:100]}...")
        else:
            print("❌ 混合检索失败")
            
    except Exception as e:
        print(f"❌ 混合检索测试失败: {e}")
        import traceback
        traceback.print_exc()

def test_vector_retrieval():
    """测试向量检索功能"""
    print("\n" + "=" * 60)
    print("测试向量检索功能")
    print("=" * 60)
    
    try:
        # 初始化RAG系统
        print("初始化RAG系统...")
        rag = RAGSystem()
        
        # 测试向量检索
        print("\n测试向量检索...")
        result = rag.retrieve("ai", mode="vector")
        
        if result:
            print(f"✅ 向量检索成功")
            print(f"检索到 {len(result.get('documents', []))} 个文档片段")
            
            # 显示前3个结果
            for i, doc in enumerate(result.get('documents', [])[:3]):
                print(f"\n结果 {i+1}:")
                print(f"  分数: {doc.get('score', 0):.4f}")
                print(f"  内容: {doc.get('text', '')[:100]}...")
        else:
            print("❌ 向量检索失败")
            
    except Exception as e:
        print(f"❌ 向量检索测试失败: {e}")
        import traceback
        traceback.print_exc()

def test_cli_commands():
    """测试CLI命令"""
    print("\n" + "=" * 60)
    print("测试CLI命令")
    print("=" * 60)
    
    commands = [
        "python cli.py retrieve -q ai --mode bm25",
        "python cli.py retrieve -q ai --mode hybrid",
        "python cli.py retrieve -q ai --mode vector",
    ]
    
    for cmd in commands:
        print(f"\n执行命令: {cmd}")
        os.system(cmd)

if __name__ == "__main__":
    print("开始测试BM25和混合检索功能...")
    
    # 测试各种检索模式
    test_vector_retrieval()
    test_bm25_retrieval()
    test_hybrid_retrieval()
    
    # 测试CLI命令
    test_cli_commands()
    
    print("\n" + "=" * 60)
    print("测试完成")
    print("=" * 60)