#!/usr/bin/env python3
"""
批量更新提示词文件，同步到20个阶段
"""

import re
from pathlib import Path

PROMPTS_DIR = Path("/home/xia-lab/wutx/00test/bbav5/prompts")

# 阶段名称映射（旧 -> 新）
STAGE_MAPPINGS = {
    # 核心阶段重命名
    '`identify`': '`gene-identify`',
    '`motif`': '`motif-analysis`',
    '`domain`': '`domain-analysis`',
    '`promoter`': '`promoter-analysis`',
    
    # 确保新阶段名称一致性
    '阶段1：基因家族鉴定': '阶段1：基因家族鉴定 (`gene-identify`)',
    '阶段2：系统发育分析': '阶段2：系统发育分析 (`phylogeny`)',
    '阶段3：保守基序分析': '阶段3：保守基序分析 (`motif-analysis`)',
    '阶段4：结构域分析': '阶段4：结构域分析 (`domain-analysis`)',
    '阶段5：蛋白结构分析': '阶段5：蛋白结构分析 (`structure`)',
    '阶段6：启动子分析': '阶段6：启动子分析 (`promoter-analysis`)',
    '阶段7：共线性分析': '阶段7：共线性分析 (`synteny`)',
    '阶段8：染色体定位': '阶段8：染色体定位 (`karyotype`)',
    '阶段9：蛋白特性分析': '阶段9：蛋白特性分析 (`pep`)',
    '阶段10：表达模式分析': '阶段10：表达模式分析 (`expression`)',
    '阶段11：基因注释': '阶段11：基因注释 (`annotation`)',
    
    # 新增阶段
    '阶段12：Ka/Ks分析': '阶段12：Ka/Ks分析 (`kaks`)',
    '阶段13：种间分析': '阶段13：种间分析 (`interspecies`)',
    '阶段14：蛋白互作分析': '阶段14：蛋白互作分析 (`ppi`)',
    '阶段15：基因家族汇总': '阶段15：基因家族汇总 (`summary`)',
    '阶段16：RNA-seq质控': '阶段16：RNA-seq质控 (`rnaseq-qc`)',
    '阶段17：RNA-seq比对': '阶段17：RNA-seq比对 (`rnaseq-align`)',
    '阶段18：RNA-seq定量': '阶段18：RNA-seq定量 (`rnaseq-quant`)',
    '阶段19：差异表达分析': '阶段19：差异表达分析 (`rnaseq-de`)',
    '阶段20：富集分析': '阶段20：富集分析 (`rnaseq-enrich`)',
}

# 需要更新的文件
FILES_TO_UPDATE = [
    'execution_agent.md',
    'result_parser_agent.md',
    'validator_agent.md',
    'diagnosis_agent.md',
]

def update_file(filepath: Path) -> dict:
    """更新单个文件"""
    content = filepath.read_text(encoding='utf-8')
    original_content = content
    
    changes = []
    
    # 应用映射
    for old, new in STAGE_MAPPINGS.items():
        if old in content:
            count = content.count(old)
            content = content.replace(old, new)
            changes.append(f"'{old}' -> '{new}' ({count}处)")
    
    # 检查是否需要更新阶段数量说明
    if '11个阶段' in content:
        content = content.replace('11个阶段', '20个阶段')
        changes.append("'11个阶段' -> '20个阶段'")
    
    if '标准分析流程（11个阶段）' in content:
        content = content.replace('标准分析流程（11个阶段）', '标准分析流程（20个阶段）')
        changes.append("更新标题")
    
    # 保存文件
    if content != original_content:
        filepath.write_text(content, encoding='utf-8')
        return {
            'file': filepath.name,
            'updated': True,
            'changes': changes
        }
    else:
        return {
            'file': filepath.name,
            'updated': False,
            'changes': []
        }

def main():
    print("=" * 60)
    print("批量更新提示词文件")
    print("=" * 60)
    print()
    
    results = []
    for filename in FILES_TO_UPDATE:
        filepath = PROMPTS_DIR / filename
        if filepath.exists():
            result = update_file(filepath)
            results.append(result)
            
            status = "✅ 已更新" if result['updated'] else "⏭️ 无需更新"
            print(f"{status}: {filename}")
            
            if result['changes']:
                for change in result['changes']:
                    print(f"    - {change}")
        else:
            print(f"❌ 文件不存在: {filename}")
    
    print()
    print("=" * 60)
    print("更新汇总")
    print("=" * 60)
    
    updated_count = sum(1 for r in results if r['updated'])
    print(f"已更新文件: {updated_count}/{len(FILES_TO_UPDATE)}")
    print()
    
    for r in results:
        if r['updated']:
            print(f"✅ {r['file']}")
            for change in r['changes']:
                print(f"   - {change}")
    
    print("=" * 60)

if __name__ == "__main__":
    main()