"""
知识库检索性能分析脚本
详细分析每个环节的耗时
"""

import asyncio
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from app.services.mid.kb_search_dynamic_tool import KbSearchDynamicTool, KbSearchDynamicConfig
from app.services.retrieval.vector_retriever import VectorRetriever
from app.core.config import get_settings
from app.core.qdrant_client import QdrantClient


async def profile_embedding_generation(query: str):
    """分析 embedding 生成耗时"""
    from app.services.embedding import get_embedding_provider
    
    start = time.time()
    embedding_service = await get_embedding_provider()
    init_time = time.time() - start
    
    start = time.time()
    embedding = await embedding_service.embed_query(query)
    embed_time = time.time() - start
    
    # 获取 embedding 向量（兼容不同 provider）
    if hasattr(embedding, 'embedding_full'):
        vector = embedding.embedding_full
    elif hasattr(embedding, 'embedding'):
        vector = embedding.embedding
    else:
        vector = embedding
    
    return {
        "init_time_ms": init_time * 1000,
        "embed_time_ms": embed_time * 1000,
        "dimension": len(vector),
    }


async def profile_qdrant_search(tenant_id: str, query_vector: list, metadata_filter: dict = None):
    """分析 Qdrant 搜索耗时"""
    from app.core.qdrant_client import get_qdrant_client
    
    client = await get_qdrant_client()
    
    # 获取 collections
    start = time.time()
    qdrant_client = await client.get_client()
    collections = await qdrant_client.get_collections()
    safe_tenant_id = tenant_id.replace('@', '_')
    prefix = f"kb_{safe_tenant_id}"
    tenant_collections = [
        c.name for c in collections.collections
        if c.name.startswith(prefix)
    ]
    list_collections_time = time.time() - start
    
    # 逐个 collection 搜索
    collection_times = []
    for collection_name in tenant_collections:
        start = time.time()
        exists = await qdrant_client.collection_exists(collection_name)
        check_time = time.time() - start
        
        if not exists:
            collection_times.append({
                "collection": collection_name,
                "exists": False,
                "time_ms": check_time * 1000,
            })
            continue
        
        start = time.time()
        # 构建 filter
        qdrant_filter = None
        if metadata_filter:
            from qdrant_client.models import FieldCondition, Filter, MatchValue
            must_conditions = []
            for key, value in metadata_filter.items():
                field_path = f"metadata.{key}"
                condition = FieldCondition(
                    key=field_path,
                    match=MatchValue(value=value),
                )
                must_conditions.append(condition)
            qdrant_filter = Filter(must=must_conditions) if must_conditions else None
        
        try:
            results = await qdrant_client.query_points(
                collection_name=collection_name,
                query=query_vector,
                using="full",  # 使用 full 向量
                limit=5,
                score_threshold=0.5,
                query_filter=qdrant_filter,
            )
        except Exception as e:
            if "vector name" in str(e).lower():
                # 尝试不使用 vector name
                results = await qdrant_client.query_points(
                    collection_name=collection_name,
                    query=query_vector,
                    limit=5,
                    score_threshold=0.5,
                    query_filter=qdrant_filter,
                )
            else:
                raise
        search_time = time.time() - start
        
        collection_times.append({
            "collection": collection_name,
            "exists": True,
            "check_time_ms": check_time * 1000,
            "search_time_ms": search_time * 1000,
            "results_count": len(results.points),
        })
    
    return {
        "list_collections_time_ms": list_collections_time * 1000,
        "collections_count": len(tenant_collections),
        "collection_times": collection_times,
    }


async def profile_full_kb_search():
    """分析完整的知识库搜索流程"""
    settings = get_settings()
    
    print("=" * 80)
    print("知识库检索性能分析")
    print("=" * 80)
    
    # 1. 分析 Embedding 生成
    print("\n📊 1. Embedding 生成分析")
    print("-" * 80)
    query = "三年级语文学习"
    embed_result = await profile_embedding_generation(query)
    print(f"  初始化时间: {embed_result['init_time_ms']:.2f} ms")
    print(f"  Embedding 生成时间: {embed_result['embed_time_ms']:.2f} ms")
    print(f"  向量维度: {embed_result['dimension']}")
    
    # 2. 分析 Qdrant 搜索
    print("\n📊 2. Qdrant 搜索分析")
    print("-" * 80)
    
    # 先生成 embedding
    from app.services.embedding import get_embedding_provider
    embedding_service = await get_embedding_provider()
    embedding_result = await embedding_service.embed_query(query)
    # 获取 embedding 向量（兼容不同 provider）
    if hasattr(embedding_result, 'embedding_full'):
        query_vector = embedding_result.embedding_full
    elif hasattr(embedding_result, 'embedding'):
        query_vector = embedding_result.embedding
    else:
        query_vector = embedding_result
    
    tenant_id = "szmp@ash@2026"
    metadata_filter = {"grade": "三年级", "subject": "语文"}
    
    qdrant_result = await profile_qdrant_search(tenant_id, query_vector, metadata_filter)
    print(f"  获取 collections 列表时间: {qdrant_result['list_collections_time_ms']:.2f} ms")
    print(f"  Collections 数量: {qdrant_result['collections_count']}")
    print(f"\n  各 Collection 搜索耗时:")
    for ct in qdrant_result['collection_times']:
        if ct['exists']:
            print(f"    - {ct['collection']}: {ct['search_time_ms']:.2f} ms (结果: {ct['results_count']} 条)")
        else:
            print(f"    - {ct['collection']}: 不存在 ({ct['time_ms']:.2f} ms)")
    
    total_search_time = sum(
        ct.get('search_time_ms', 0) for ct in qdrant_result['collection_times']
    )
    print(f"\n  总搜索时间（串行）: {total_search_time:.2f} ms")
    
    # 3. 分析完整流程
    print("\n📊 3. 完整 KB Search 流程分析")
    print("-" * 80)
    
    engine = create_async_engine(settings.database_url)
    async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
    
    async with async_session() as session:
        config = KbSearchDynamicConfig(
            enabled=True,
            top_k=5,
            timeout_ms=30000,  # 30秒
            min_score_threshold=0.5,
        )
        
        tool = KbSearchDynamicTool(session=session, config=config)
        
        # 记录各阶段时间
        stages = []
        
        start_total = time.time()
        
        # 执行搜索
        start = time.time()
        result = await tool.execute(
            query=query,
            tenant_id=tenant_id,
            scene="学习方案",
            top_k=5,
            context=metadata_filter,
        )
        total_time = (time.time() - start_total) * 1000
        
        print(f"  总耗时: {total_time:.2f} ms")
        print(f"  结果: success={result.success}, hits={len(result.hits)}")
        print(f"  工具内部耗时: {result.duration_ms} ms")
        
        # 计算时间差（工具内部 vs 外部测量）
        overhead = total_time - result.duration_ms
        print(f"  额外开销（初始化等）: {overhead:.2f} ms")
    
    # 4. 性能瓶颈分析
    print("\n📊 4. 性能瓶颈分析")
    print("-" * 80)
    
    embedding_time = embed_result['embed_time_ms']
    qdrant_time = total_search_time
    total_measured = embedding_time + qdrant_time
    
    print(f"  Embedding 生成: {embedding_time:.2f} ms ({embedding_time/total_measured*100:.1f}%)")
    print(f"  Qdrant 搜索: {qdrant_time:.2f} ms ({qdrant_time/total_measured*100:.1f}%)")
    print(f"  其他开销: {total_time - total_measured:.2f} ms")
    
    print("\n" + "=" * 80)
    print("优化建议:")
    print("=" * 80)
    
    if embedding_time > 1000:
        print("  ⚠️ Embedding 生成较慢，考虑:")
        print("     - 使用更快的 embedding 模型")
        print("     - 增加 embedding 服务缓存")
    
    if qdrant_time > 1000:
        print("  ⚠️ Qdrant 搜索较慢，考虑:")
        print("     - 并行查询多个 collections")
        print("     - 优化 Qdrant 索引")
        print("     - 减少 collections 数量")
    
    if len(qdrant_result['collection_times']) > 3:
        print(f"  ⚠️ Collections 数量较多 ({len(qdrant_result['collection_times'])} 个)")
        print("     - 建议合并或归档空/少数据的 collections")


if __name__ == "__main__":
    asyncio.run(profile_full_kb_search())
