"""
RAG-MCP 增强功能测试脚本
测试新增功能：相似度过滤、元数据过滤、增量索引管理、批量查询
"""
import sys
import os
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from loguru import logger
from llama_index.core import Document

from src.core import RAGSystem
from src.config import get_config_manager


def print_section(title: str):
    """打印分节标题"""
    print("\n" + "=" * 80)
    print(f"  {title}")
    print("=" * 80 + "\n")


def test_similarity_postprocessor():
    """测试1: 相似度过滤（SimilarityPostprocessor）"""
    print_section("测试 1: 相似度过滤 (SimilarityPostprocessor)")
    
    try:
        # 临时启用相似度过滤
        config = get_config_manager()
        original_filter = config.retrieval.use_similarity_filter
        original_cutoff = config.retrieval.similarity_cutoff
        
        config.retrieval.use_similarity_filter = True
        config.retrieval.similarity_cutoff = 0.3
        
        print(f"✓ 配置: use_similarity_filter={config.retrieval.use_similarity_filter}")
        print(f"✓ 配置: similarity_cutoff={config.retrieval.similarity_cutoff}")
        
        # 创建RAG系统
        rag = RAGSystem(config=config)
        
        # 检索测试
        print("\n执行检索测试...")
        result = rag.retrieve(
            question="什么是人工智能？",
            mode="hybrid",
            top_k=10
        )
        
        if result["success"]:
            print(f"✅ 相似度过滤测试通过")
            print(f"   检索到 {result['count']} 个文档（过滤后）")
            if result["results"]:
                print(f"   最高分数: {result['results'][0].get('score', 'N/A')}")
        else:
            print(f"❌ 测试失败: {result.get('error')}")
        
        # 恢复原配置
        config.retrieval.use_similarity_filter = original_filter
        config.retrieval.similarity_cutoff = original_cutoff
        
        return result["success"]
        
    except Exception as e:
        print(f"❌ 测试异常: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_metadata_filters():
    """测试2: 元数据过滤"""
    print_section("测试 2: 元数据过滤 (MetadataFilters)")
    
    try:
        rag = RAGSystem()
        
        # 测试元数据过滤
        print("执行元数据过滤测试...")
        
        metadata_filters = {
            "file_name": {
                "operator": "TEXT_MATCH",
                "value": ".txt"
            }
        }
        
        print(f"✓ 过滤条件: file_name TEXT_MATCH '.txt'")
        
        result = rag.retrieve(
            question="机器学习",
            mode="hybrid",
            top_k=5,
            metadata_filters=metadata_filters
        )
        
        if result["success"]:
            print(f"✅ 元数据过滤测试通过")
            print(f"   检索到 {result['count']} 个文档")
            for idx, doc in enumerate(result["results"][:3], 1):
                file_name = doc.get("metadata", {}).get("file_name", "未知")
                print(f"   {idx}. 文件: {file_name}")
        else:
            print(f"❌ 测试失败: {result.get('error')}")
        
        return result["success"]
        
    except Exception as e:
        print(f"❌ 测试异常: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_incremental_index_management():
    """测试3: 增量索引管理"""
    print_section("测试 3: 增量索引管理")
    
    try:
        rag = RAGSystem()
        
        # 3.1 列出现有文档
        print("3.1 列出现有文档...")
        result = rag.list_documents()
        
        if result["success"]:
            total_docs = result.get("total_count", 0)
            print(f"✓ 当前索引中有 {total_docs} 个文档")
        else:
            print(f"⚠️  无法列出文档: {result.get('error')}")
            return False
        
        # 3.2 插入新文档
        print("\n3.2 插入新文档...")
        new_docs = [
            Document(
                text="这是一个测试文档，用于验证增量索引功能。包含关键词：测试、增量、索引。",
                metadata={
                    "source": "test",
                    "category": "enhancement",
                    "test_id": "test_doc_001"
                }
            )
        ]
        
        result = rag.insert_documents(new_docs)
        
        if result["success"]:
            print(f"✅ 插入文档成功")
            print(f"   插入数量: {result.get('inserted_count', 0)}")
            inserted_doc_id = None
        else:
            print(f"❌ 插入失败: {result.get('error')}")
            return False
        
        # 3.3 再次列出文档，验证插入
        print("\n3.3 验证文档已插入...")
        result = rag.list_documents()
        
        if result["success"]:
            new_total = result.get("total_count", 0)
            print(f"✓ 当前文档数: {new_total}")
            
            # 查找我们刚插入的文档
            for doc in result.get("documents", []):
                if doc.get("metadata", {}).get("test_id") == "test_doc_001":
                    inserted_doc_id = doc.get("doc_id")
                    print(f"✓ 找到插入的文档, ID: {inserted_doc_id}")
                    break
        
        # 3.4 获取文档详情
        if inserted_doc_id:
            print(f"\n3.4 获取文档详情...")
            result = rag.get_document(inserted_doc_id)
            
            if result["success"]:
                print(f"✅ 获取文档成功")
                print(f"   文档ID: {result.get('doc_id')}")
                print(f"   文本长度: {len(result.get('text', ''))}")
                print(f"   元数据: {result.get('metadata')}")
            else:
                print(f"❌ 获取文档失败: {result.get('error')}")
        
        # 3.5 删除测试文档
        if inserted_doc_id:
            print(f"\n3.5 删除测试文档...")
            result = rag.delete_document(inserted_doc_id)
            
            if result["success"]:
                print(f"✅ 删除文档成功")
                print(f"   删除的文档ID: {result.get('deleted_doc_id')}")
            else:
                print(f"❌ 删除失败: {result.get('error')}")
                return False
        
        # 3.6 验证删除
        print("\n3.6 验证文档已删除...")
        result = rag.list_documents()
        
        if result["success"]:
            final_total = result.get("total_count", 0)
            print(f"✓ 最终文档数: {final_total}")
            
            if final_total == total_docs:
                print(f"✅ 增量索引管理测试通过 - 文档数恢复到初始状态")
                return True
            else:
                print(f"⚠️  文档数不匹配: 期望 {total_docs}, 实际 {final_total}")
                return False
        
        return True
        
    except Exception as e:
        print(f"❌ 测试异常: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_batch_query():
    """测试4: 批量异步查询"""
    print_section("测试 4: 批量异步查询")
    
    try:
        rag = RAGSystem()
        
        # 准备批量问题
        questions = [
            "什么是人工智能？",
            "什么是机器学习？",
            "什么是深度学习？"
        ]
        
        print(f"准备批量查询 {len(questions)} 个问题...")
        for idx, q in enumerate(questions, 1):
            print(f"  {idx}. {q}")
        
        # 执行批量查询
        print("\n执行批量查询...")
        results = rag.batch_query(
            questions=questions,
            mode="hybrid",
            top_k=3
        )
        
        if isinstance(results, list) and len(results) > 0:
            print(f"✅ 批量查询测试通过")
            print(f"   返回结果数: {len(results)}")
            
            # 显示每个结果的摘要
            for idx, result in enumerate(results, 1):
                success = result.get("success", False)
                question = result.get("question", "")
                answer_len = len(result.get("answer", ""))
                sources_count = len(result.get("sources", []))
                
                status = "✓" if success else "✗"
                print(f"   {status} 问题{idx}: {question[:30]}...")
                print(f"      答案长度: {answer_len}, 来源数: {sources_count}")
            
            return True
        else:
            print(f"❌ 测试失败: 返回结果格式不正确")
            return False
        
    except Exception as e:
        print(f"❌ 测试异常: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """主测试函数"""
    print("\n" + "=" * 80)
    print("  RAG-MCP 增强功能测试")
    print("  测试范围: 相似度过滤、元数据过滤、增量索引管理、批量查询")
    print("=" * 80)
    
    # 检查索引是否存在
    config = get_config_manager()
    rag = RAGSystem(config=config)
    
    print("\n检查索引状态...")
    stats = rag.get_index_stats()
    
    if not stats.get("vector_index_exists"):
        print("\n❌ 错误: 索引不存在，请先构建索引")
        print("   运行: python cli.py index --dir ./my_knowledge_base")
        return False
    
    print("✓ 索引已存在，可以进行测试\n")
    
    # 运行测试
    results = {}
    
    # 测试1: 相似度过滤
    results["similarity_filter"] = test_similarity_postprocessor()
    
    # 测试2: 元数据过滤
    results["metadata_filter"] = test_metadata_filters()
    
    # 测试3: 增量索引管理
    results["incremental_index"] = test_incremental_index_management()
    
    # 测试4: 批量查询
    results["batch_query"] = test_batch_query()
    
    # 总结
    print_section("测试总结")
    
    total_tests = len(results)
    passed_tests = sum(1 for v in results.values() if v)
    
    print(f"总测试数: {total_tests}")
    print(f"通过数量: {passed_tests}")
    print(f"失败数量: {total_tests - passed_tests}")
    print(f"\n通过率: {passed_tests / total_tests * 100:.1f}%\n")
    
    # 详细结果
    for test_name, passed in results.items():
        status = "✅ 通过" if passed else "❌ 失败"
        print(f"  {test_name:20s} - {status}")
    
    print("\n" + "=" * 80)
    
    if passed_tests == total_tests:
        print("🎉 所有测试通过！增强功能工作正常。")
        return True
    else:
        print("⚠️  部分测试失败，请检查日志。")
        return False


if __name__ == "__main__":
    try:
        success = main()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n👋 测试已取消")
        sys.exit(130)
    except Exception as e:
        print(f"\n❌ 测试程序异常: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
