#!/usr/bin/env python3
"""
pipeline_backend.py 可复现性模块集成补丁

使用方法:
1. 导入 ReproducibilityBundle
2. 在 run_stage 函数开头初始化 bundle
3. 在命令执行后记录命令
4. 在函数返回前生成可复现性文件
"""

# ============ 导入部分（添加到文件顶部）============
IMPORT_PATCH = '''
# 可复现性支持
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))
from gfa_core.reproducibility import ReproducibilityBundle
'''

# ============ run_stage 函数修改 ============
RUN_STAGE_PATCH = '''
    # 初始化可复现性 bundle
    reproducibility_bundle = None
    try:
        reproducibility_bundle = ReproducibilityBundle(
            job_id=job.job_id,
            prefix=prefix,
            workspace_root=str(WORKSPACE_ROOT)
        )
    except Exception as e:
        # 如果初始化失败，不影响主流程
        print(f"[WARN] ReproducibilityBundle init failed: {e}")
'''

# ============ 在每个 stage 命令构建后添加 ============
RECORD_CMD_PATCH = '''
        # 记录命令到可复现性 bundle
        if reproducibility_bundle:
            try:
                reproducibility_bundle.add_command(
                    stage=stage,
                    cmd=cmd,
                    description=f"Run {stage} analysis"
                )
            except Exception as e:
                print(f"[WARN] Failed to record command: {e}")
'''

# ============ 在函数返回前添加 ============
GENERATE_REPRO_PATCH = '''
    # 生成可复现性文件
    if reproducibility_bundle:
        try:
            reproducibility_bundle.generate_commands_sh()
            reproducibility_bundle.generate_environment_yml()
            reproducibility_bundle.generate_checksums()
            reproducibility_bundle.copy_to_project_reproducibility()
            print(f"[INFO] Reproducibility files generated for {stage}")
        except Exception as e:
            print(f"[WARN] Failed to generate reproducibility files: {e}")
'''


def apply_patch_to_pipeline_backend():
    """
    自动应用补丁到 pipeline_backend.py
    """
    import re
    
    file_path = Path("/home/xia-lab/wutx/00test/bbav5/apps/web/pipeline_backend.py")
    content = file_path.read_text()
    
    # 1. 添加导入
    if "from gfa_core.reproducibility import ReproducibilityBundle" not in content:
        # 在文件开头添加导入
        import_section = """# web/pipeline_backend.py
import json
import os
import re
import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# 可复现性支持
sys.path.insert(0, str(Path(__file__).parent.parent))
from gfa_core.reproducibility import ReproducibilityBundle
"""
        content = content.replace(
            "# web/pipeline_backend.py\nimport json",
            import_section
        )
        print("✅ Added ReproducibilityBundle import")
    
    # 2. 在 run_stage 函数中添加初始化
    if "reproducibility_bundle = None" not in content:
        # 找到 run_stage 函数中 job: JobInfo 后的位置
        pattern = r"(def run_stage\([^)]+job: JobInfo[^)]+\):.*?ok, msg = can_start_new_task\(\).*?if not ok:.*?return msg, None)"
        
        def add_init(match):
            return match.group(1) + """

    # 初始化可复现性 bundle
    reproducibility_bundle = None
    try:
        reproducibility_bundle = ReproducibilityBundle(
            job_id=job.job_id,
            prefix=prefix,
            workspace_root=str(WORKSPACE_ROOT)
        )
    except Exception as e:
        # 如果初始化失败，不影响主流程
        print(f"[WARN] ReproducibilityBundle init failed: {e}")"""
        
        content = re.sub(pattern, add_init, content, flags=re.DOTALL)
        print("✅ Added ReproducibilityBundle initialization")
    
    # 保存修改
    file_path.write_text(content)
    print(f"✅ Patch applied to {file_path}")


if __name__ == "__main__":
    apply_patch_to_pipeline_backend()