#!/usr/bin/env python3
"""
Annotation Skill 使用示例
展示如何在 pipeline 中集成 eggNOG-mapper 功能注释
"""

import subprocess
from pathlib import Path


def run_annotation_skill(protein_fasta: Path, output_dir: Path, **kwargs):
    """
    调用 annotation skill 进行功能注释
    
    Args:
        protein_fasta: 蛋白序列FASTA文件路径
        output_dir: 输出目录
        **kwargs: 可选参数 (cpu, mode, evalue, etc.)
    
    Returns:
        dict: 包含 annotations 文件路径和统计信息
    """
    skill_path = Path(__file__).parent / "run.py"
    
    cmd = [
        "python3", str(skill_path),
        "--input", str(protein_fasta),
        "--output", str(output_dir),
    ]
    
    # 添加可选参数
    if "cpu" in kwargs:
        cmd.extend(["--cpu", str(kwargs["cpu"])])
    if "mode" in kwargs:
        cmd.extend(["--mode", kwargs["mode"]])
    if "evalue" in kwargs:
        cmd.extend(["--evalue", str(kwargs["evalue"])])
    if "prefix" in kwargs:
        cmd.extend(["--prefix", kwargs["prefix"]])
    
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.returncode != 0:
        raise RuntimeError(f"Annotation failed: {result.stderr}")
    
    import json
    return json.loads(result.stdout)


def integrate_with_pipeline(gene_family_pep: Path, work_dir: Path):
    """
    集成到基因家族分析 pipeline 的示例
    
    在基因鉴定后，对鉴定出的基因家族成员进行功能注释
    """
    annotation_dir = work_dir / "05_annotation"
    annotation_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"[Pipeline] Running functional annotation for {gene_family_pep.name}")
    
    # 调用 annotation skill
    result = run_annotation_skill(
        protein_fasta=gene_family_pep,
        output_dir=annotation_dir,
        prefix="gene_family",
        cpu=8,
        mode="diamond",
        evalue=0.001,
    )
    
    if result["status"] == "success":
        print(f"[Pipeline] Annotation completed!")
        print(f"  - Annotations file: {result['annotations']}")
        print(f"  - Report: {result['report']}")
        print(f"  - Stats: {result['stats']['total_genes']} genes, "
              f"{result['stats']['with_go']} with GO, "
              f"{result['stats']['with_kegg']} with KEGG")
        
        return result
    else:
        raise RuntimeError("Annotation failed")


if __name__ == "__main__":
    # 示例用法
    print("=" * 60)
    print("Annotation Skill Integration Example")
    print("=" * 60)
    
    # 创建示例蛋白文件
    demo_pep = Path("/tmp/demo_genes.pep")
    demo_pep.write_text(">Gene001\nMKTLLLPLLLLVLLLLLVQGSISAQAAAAAA\n"
                       ">Gene002\nMKTLLLPLLLLVLLLLLVQGSISAQBBBBBB\n")
    
    # 运行注释
    result = integrate_with_pipeline(demo_pep, Path("/tmp/demo_annotation"))
    
    print("\n" + "=" * 60)
    print("Result:")
    print(f"  Status: {result['status']}")
    print(f"  Output dir: {result['output_dir']}")
    print(f"  Annotations: {result['annotations']}")