"""
PubMed 文献智能分析示例（Pydantic AI 版本）

本示例展示如何使用基于 Pydantic AI 的 PaperAnalyzer 进行文献分析。

Pydantic AI 相比传统 LLM 调用的优势：
1. 类型安全：输出自动验证为 Pydantic 模型，无需手动解析
2. 依赖注入：通过 AnalysisDeps 管理配置和缓存
3. 结构化输出：直接使用强类型对象，IDE 自动补全友好
4. Token 追踪：内置用量统计，便于成本监控

学习目标：
- 配置 AnalysisConfig 启用 Pydantic AI
- 使用 PaperAnalyzer 进行单篇和批量分析
- 处理 AnalysisResult 结构化结果
- 追踪 token 使用量

迁移说明（从 LangChain 到 Pydantic AI）：

主要变化：
1. 导入方式改变：
   - 旧：from langchain_openai import ChatOpenAI
   - 新：from information_composer.pubmed.analyzer import PaperAnalyzer

2. 初始化方式改变：
   - 旧：llm = ChatOpenAI(...)
   - 新：analyzer = PaperAnalyzer(AnalysisConfig(...))

3. 调用方式改变：
   - 旧：result = await chain.ainvoke(...)
   - 新：result = await analyzer.analyze_paper(paper)

4. 结果处理：
   - 旧：需要手动解析 JSON 字符串
   - 新：自动返回 Pydantic 模型，类型安全
"""

import asyncio
import os
from pathlib import Path

from information_composer.pubmed.analyzer import (
    AnalysisConfig,
    PaperAnalyzer,
    PaperInput,
)


async def example_single_summary():
    """示例1: 单篇论文总结分析

    本示例展示如何使用 Pydantic AI 进行单篇论文的总结分析。

    Pydantic AI 优势：
    - 自动验证输出格式为 SummaryResult Pydantic 模型
    - 无需手动处理 JSON 解析错误
    - IDE 提供 result.summary.main_findings 等属性的类型提示
    """
    print("=" * 80)
    print("示例1: 单篇论文总结分析")
    print("=" * 80)

    # 配置分析器 - 只进行论文总结
    # use_pydantic_ai=True 启用 Pydantic AI Agent 实现
    # 设置为 False 则使用 LangChain 实现（向后兼容）
    config = AnalysisConfig(
        analysis_types=["summary"],
        llm_model="qwen-plus-latest",
        temperature=0.1,
        cache_enabled=True,
        use_pydantic_ai=True,  # 关键：启用 Pydantic AI
    )

    # 创建分析器
    # PaperAnalyzer 根据 use_pydantic_ai 自动选择底层实现
    analyzer = PaperAnalyzer(config)

    # 准备论文数据
    # PaperInput 是 Pydantic 模型，自动验证输入数据
    paper = PaperInput(
        pmid="12345678",
        title="CRISPR-Cas9 mediated genome editing in rice for drought tolerance",
        abstract=(
            "Drought is a major environmental stress affecting rice yield worldwide. "
            "We developed a CRISPR-Cas9 system to edit key drought-responsive genes in rice. "
            "Three genes (OsDREB1A, OsNAC5, and OsLEA3) were simultaneously edited using "
            "multiplex CRISPR technology. The edited rice lines showed 30% higher grain yield "
            "under drought stress compared to wild-type plants. Gene expression analysis revealed "
            "enhanced activation of stress-responsive pathways. This study demonstrates the "
            "potential of genome editing for developing climate-resilient crops."
        ),
    )

    # 执行分析
    print("\n正在分析论文...")
    result = await analyzer.analyze_paper(paper)

    # 显示结果
    # AnalysisResult 是 Pydantic 模型，所有字段都有类型提示
    print(f"\n论文标题: {result.title}")
    print(f"PMID: {result.pmid}")

    # result.summary 是 SummaryResult 类型
    # Pydantic AI 自动验证 LLM 输出符合 SummaryResult 结构
    if result.summary:
        print("\n主要发现:")
        for i, finding in enumerate(result.summary.main_findings, 1):
            print(f"  {i}. {finding}")

        print("\n创新点:")
        for i, innovation in enumerate(result.summary.innovations, 1):
            print(f"  {i}. {innovation}")

        print("\n核心结论:")
        print(f"  {result.summary.conclusions}")

    # 显示处理元数据
    # ProcessingMetadata 包含 LLM 使用情况
    print("\n处理元数据:")
    print(f"  模型: {result.processing_metadata.llm_model}")
    print(f"  耗时: {result.processing_metadata.processing_time:.2f} 秒")
    print(f"  缓存命中: {result.processing_metadata.cache_hit}")

    # Token 使用量追踪（Pydantic AI 原生支持）
    if result.processing_metadata.tokens_used:
        print(f"  Token 使用量: {result.processing_metadata.tokens_used}")

    return result


async def example_domain_analysis():
    """示例2: 领域判定分析

    本示例展示如何使用 Pydantic AI 进行论文领域判定。

    Pydantic AI 优势：
    - DomainResult 自动验证领域评分在 0-1 范围内
    - primary_domain 必须在 relevant_domains 列表中
    - 无需手动验证字段有效性
    """
    print("\n" + "=" * 80)
    print("示例2: 领域判定分析")
    print("=" * 80)

    # 配置分析器 - 只进行领域判定
    # domain_list 定义候选领域
    config = AnalysisConfig(
        analysis_types=["domain"],
        domain_list=[
            "Genome Editing",
            "Plant Breeding",
            "Climate Change",
            "Cancer Research",
        ],
        llm_model="qwen-plus-latest",
        cache_enabled=True,
        use_pydantic_ai=True,  # 启用 Pydantic AI
    )

    # 创建分析器
    analyzer = PaperAnalyzer(config)

    # 准备论文数据
    paper = PaperInput(
        pmid="23456789",
        title="Application of CRISPR technology in crop improvement",
        abstract=(
            "CRISPR-Cas9 has revolutionized plant breeding by enabling precise genome editing. "
            "This review discusses recent advances in CRISPR applications for crop improvement, "
            "including yield enhancement, disease resistance, and stress tolerance. "
            "We highlight successful examples in rice, wheat, and maize, and discuss "
            "regulatory challenges for genome-edited crops."
        ),
    )

    # 执行分析
    print("\n正在分析论文领域...")
    result = await analyzer.analyze_paper(paper)

    # 显示结果
    # result.domain_analysis 是 DomainResult 类型
    print(f"\n论文标题: {result.title}")

    if result.domain_analysis:
        print("\n相关领域:")
        # domain_scores 是 dict[str, float]，自动验证评分范围
        for domain in result.domain_analysis.relevant_domains:
            score = result.domain_analysis.domain_scores[domain]
            print(f"  - {domain}: {score:.2f}")

        print(f"\n主要领域: {result.domain_analysis.primary_domain}")
        print("\n判定依据:")
        print(f"  {result.domain_analysis.reasoning}")

    return result


async def example_full_analysis():
    """示例3: 完整分析（总结 + 领域判定）

    本示例展示如何同时进行多种分析。

    Pydantic AI 优势：
    - 并行执行多个 Agent（summary 和 domain）
    - 每个 Agent 的输出独立验证
    - 统一的 AnalysisResult 包含所有分析结果
    """
    print("\n" + "=" * 80)
    print("示例3: 完整分析（总结 + 领域判定）")
    print("=" * 80)

    # 配置分析器 - 包含所有分析类型
    config = AnalysisConfig(
        analysis_types=["summary", "domain"],
        domain_list=["Epigenetics", "Gene Expression", "Plant Development"],
        llm_model="qwen-plus-latest",
        cache_enabled=True,
        use_pydantic_ai=True,  # 启用 Pydantic AI
    )

    # 创建分析器
    analyzer = PaperAnalyzer(config)

    # 准备论文数据
    paper = PaperInput(
        pmid="34567890",
        title="Histone modifications regulate flowering time in Arabidopsis",
        abstract=(
            "Flowering time is controlled by complex epigenetic mechanisms. "
            "We investigated the role of histone H3 lysine 27 trimethylation (H3K27me3) "
            "in vernalization-mediated flowering. ChIP-seq analysis revealed that H3K27me3 "
            "marks are dynamically deposited on FLOWERING LOCUS C (FLC) during cold exposure. "
            "Mutations in Polycomb Repressive Complex 2 (PRC2) components abolished H3K27me3 "
            "deposition and prevented vernalization response. Our findings demonstrate that "
            "epigenetic regulation is essential for plant adaptation to seasonal changes."
        ),
    )

    # 执行分析
    print("\n正在执行完整分析...")
    result = await analyzer.analyze_paper(paper)

    # 显示总结结果
    if result.summary:
        print("\n【论文总结】")
        print("主要发现:")
        for i, finding in enumerate(result.summary.main_findings, 1):
            print(f"  {i}. {finding}")

        print("\n创新点:")
        for i, innovation in enumerate(result.summary.innovations, 1):
            print(f"  {i}. {innovation}")

    # 显示领域判定结果
    if result.domain_analysis:
        print("\n【领域判定】")
        print(f"主要领域: {result.domain_analysis.primary_domain}")
        print("相关领域评分:")
        for domain in result.domain_analysis.relevant_domains:
            score = result.domain_analysis.domain_scores[domain]
            print(f"  - {domain}: {score:.2f}")

    # 显示置信度
    print("\n【置信度评分】")
    for analysis_type, score in result.confidence_scores.items():
        print(f"  {analysis_type}: {score:.2f}")

    # 显示 token 使用量
    if result.processing_metadata.tokens_used:
        print(f"\n【Token 使用量】")
        print(f"  总计: {result.processing_metadata.tokens_used}")

    return result


async def example_batch_analysis():
    """示例4: 批量分析

    本示例展示如何批量分析多篇论文。

    Pydantic AI 优势：
    - 自动并发控制（通过 max_concurrent 参数）
    - 每个结果独立验证
    - 支持进度回调
    """
    print("\n" + "=" * 80)
    print("示例4: 批量论文分析")
    print("=" * 80)

    # 配置分析器
    # batch_size 和 max_concurrent 控制并发性能
    config = AnalysisConfig(
        analysis_types=["summary", "domain"],
        domain_list=["Genome Editing", "Plant Genetics", "Molecular Biology"],
        llm_model="qwen-plus-latest",
        batch_size=2,  # 每批处理2篇
        max_concurrent=2,  # 最大并发数
        cache_enabled=True,
        use_pydantic_ai=True,  # 启用 Pydantic AI
    )

    # 创建分析器
    analyzer = PaperAnalyzer(config)

    # 准备多篇论文
    papers = [
        PaperInput(
            pmid="11111111",
            title="CRISPR base editing in plants",
            abstract="We developed a cytosine base editor for precise point mutations in rice without double-strand breaks.",
        ),
        PaperInput(
            pmid="22222222",
            title="Prime editing technology for crop improvement",
            abstract="Prime editing enables insertions, deletions, and base conversions in plant genomes with high precision.",
        ),
        PaperInput(
            pmid="33333333",
            title="RNA interference in plant disease resistance",
            abstract="RNAi-based approach provides effective resistance against viral and fungal pathogens in crops.",
        ),
    ]

    # 定义进度回调
    def progress_callback(current, total):
        print(f"  进度: {current}/{total} ({current * 100 // total}%)")

    # 执行批量分析
    print(f"\n开始批量分析 {len(papers)} 篇论文...")
    results = await analyzer.analyze_batch(papers, progress_callback=progress_callback)

    # 显示结果摘要
    print(f"\n批量分析完成，成功分析 {len(results)} 篇论文\n")
    for i, result in enumerate(results, 1):
        print(f"{i}. {result.title}")
        print(f"   PMID: {result.pmid}")
        if result.domain_analysis:
            print(f"   主要领域: {result.domain_analysis.primary_domain}")
        if result.summary:
            print(f"   主要发现数: {len(result.summary.main_findings)}")
        print()

    # 导出结果
    output_dir = Path("./output")
    output_dir.mkdir(exist_ok=True)
    output_file = output_dir / "batch_analysis_results.json"

    print(f"导出结果到: {output_file}")
    analyzer.export_results(results, str(output_file))

    return results


async def example_cache_management():
    """示例5: 缓存管理

    本示例展示如何查看缓存统计信息。

    Pydantic AI 缓存：
    - 自动缓存 LLM 调用结果
    - 基于输入内容计算缓存键
    - 减少重复 API 调用，节省成本
    """
    print("\n" + "=" * 80)
    print("示例5: 缓存管理")
    print("=" * 80)

    # 创建分析器
    config = AnalysisConfig(
        analysis_types=["summary"],
        cache_enabled=True,
        use_pydantic_ai=True,  # 启用 Pydantic AI
    )
    analyzer = PaperAnalyzer(config)

    # 获取缓存统计
    stats = analyzer.get_cache_stats()
    print("\n缓存统计信息:")
    print(f"  启用状态: {stats.get('enabled', False)}")
    print(f"  缓存目录: {stats.get('cache_dir', 'N/A')}")
    print(f"  缓存文件数: {stats.get('total_files', 0)}")
    print(f"  总大小: {stats.get('total_size_mb', 0)} MB")

    return stats


async def example_token_tracking():
    """示例6: Token 使用量追踪

    本示例展示如何使用 Pydantic AI 的 token 使用量追踪功能。

    Pydantic AI 优势：
    - 自动追踪每个请求的 token 使用量
    - 支持成本计算和预算控制
    - 无需手动解析响应头
    """
    print("\n" + "=" * 80)
    print("示例6: Token 使用量追踪")
    print("=" * 80)

    # 配置分析器
    config = AnalysisConfig(
        analysis_types=["summary"],
        llm_model="qwen-plus-latest",
        cache_enabled=True,
        use_pydantic_ai=True,  # 关键：Pydantic AI 原生支持 token 追踪
    )
    analyzer = PaperAnalyzer(config)

    # 准备论文数据
    paper = PaperInput(
        pmid="44444444",
        title="Gene editing for improved crop yield",
        abstract=(
            "We used CRISPR-Cas9 to edit yield-related genes in wheat. "
            "The edited lines showed 20% higher grain yield in field trials."
        ),
    )

    # 执行分析
    print("\n正在分析论文...")
    result = await analyzer.analyze_paper(paper)

    # 显示 token 使用信息
    print(f"\n论文标题: {result.title}")
    print(f"PMID: {result.pmid}")

    # Pydantic AI 自动追踪 token 使用
    if result.processing_metadata.tokens_used:
        print(f"\nToken 使用量:")
        print(f"  总计: {result.processing_metadata.tokens_used}")
        # 可以基于 token 数量计算成本
        # 例如：假设 $0.002 / 1K tokens
        estimated_cost = (result.processing_metadata.tokens_used / 1000) * 0.002
        print(f"  估算成本: ${estimated_cost:.4f}")
    else:
        print("\nToken 使用信息不可用（可能命中缓存）")

    return result


async def main():
    """主函数 - 运行所有示例"""
    print("\n" + "=" * 80)
    print("PubMed 文献智能分析示例（Pydantic AI 版本）")
    print("=" * 80)

    # 检查环境变量
    if not os.getenv("DASHSCOPE_API_KEY"):
        print("\n错误: 未设置 DASHSCOPE_API_KEY 环境变量")
        print("请设置环境变量: export DASHSCOPE_API_KEY='your-api-key'")
        return

    try:
        # 运行示例 1: 论文总结
        await example_single_summary()

        # 运行示例 2: 领域判定
        await example_domain_analysis()

        # 运行示例 3: 完整分析
        await example_full_analysis()

        # 运行示例 4: 批量分析
        await example_batch_analysis()

        # 运行示例 5: 缓存管理
        await example_cache_management()

        # 运行示例 6: Token 追踪
        await example_token_tracking()

        print("\n" + "=" * 80)
        print("所有示例运行完成！")
        print("=" * 80)
        print("\nPydantic AI 优势总结:")
        print("  - 类型安全: 自动验证输出格式")
        print("  - 依赖注入: 通过 AnalysisDeps 管理依赖")
        print("  - Token 追踪: 原生支持使用量统计")
        print("  - IDE 友好: 完整的类型提示和自动补全")

    except Exception as e:
        print(f"\n错误: {e}")
        import traceback

        traceback.print_exc()


if __name__ == "__main__":
    # 运行主函数
    asyncio.run(main())
