#!/usr/bin/env python3
"""
批量为训练脚本添加 AUTO EVAL 功能

Usage:
    python chengyue_script/dvlm/training/add_auto_eval.py --dry-run
    python chengyue_script/dvlm/training/add_auto_eval.py
"""

import argparse
import os
import re
from pathlib import Path

AUTO_EVAL_BLOCK = '''
    # ============================================
    # AUTO EVAL: 训练成功后自动提交评测任务
    # ============================================
    echo ""
    echo "=========================================="
    echo "🚀 Auto-submitting evaluation jobs..."
    echo "=========================================="
    
    if [ -f "${PROJECT_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" ]; then
        bash "${PROJECT_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" \\
            "${OUTPUT_DIR}" \\
            --skip-existing 1 \\
            2>&1 | tee "${LOG_DIR}/auto_eval_submit.log"
        echo "✅ Evaluation jobs submitted. Check: ${OUTPUT_DIR}/eval_vlmeval/"
    else
        echo "⚠️  submit_batch_eval.sh not found, skipping auto eval"
    fi
'''


def find_training_scripts(base_dir: Path) -> list:
    """Find all training shell scripts (any .sh file that contains exit_code check)."""
    scripts = []
    for sh_file in base_dir.rglob("*.sh"):
        # Check if file contains training-related patterns
        try:
            content = sh_file.read_text()
            # Look for patterns that indicate this is a training script
            if ("exit_code" in content.lower() or "EXIT_CODE" in content) and \
               ("Success!" in content or "成功" in content):
                scripts.append(sh_file)
        except Exception:
            pass
    return sorted(scripts)


def has_auto_eval(content: str) -> bool:
    """Check if script already has auto eval."""
    return "AUTO EVAL" in content or "submit_batch_eval" in content


def add_auto_eval(content: str) -> tuple:
    """
    Add auto eval block after the success message.
    
    Returns (new_content, success)
    """
    # Pattern variants to match different script styles
    patterns = [
        # Pattern 1: if [ $exit_code -eq 0 ]; then\n    echo "✅ ... Success!"
        r'(if \[ \$exit_code -eq 0 \]; then\s*\n\s*echo "[^"]*Success!")',
        # Pattern 2: EXIT_CODE variant
        r'(if \[ \$EXIT_CODE -eq 0 \]; then\s*\n\s*echo "[^"]*Success!")',
        # Pattern 3: if [ "$EXIT_CODE" -eq 0 ]
        r'(if \[ "\$EXIT_CODE" -eq 0 \]; then\s*\n\s*echo "[^"]*Success!")',
        # Pattern 4: 成功 (Chinese)
        r'(if \[ \$exit_code -eq 0 \]; then\s*\n\s*echo "[^"]*成功[^"]*")',
        # Pattern 5: EXIT_CODE with 成功
        r'(if \[ \$EXIT_CODE -eq 0 \]; then\s*\n\s*echo "[^"]*成功[^"]*")',
    ]
    
    match = None
    for pattern in patterns:
        match = re.search(pattern, content, re.IGNORECASE)
        if match:
            break
    
    if not match:
        return content, False
    
    # Insert auto eval block after the success message
    insert_pos = match.end()
    new_content = content[:insert_pos] + AUTO_EVAL_BLOCK + content[insert_pos:]
    
    return new_content, True


def main():
    parser = argparse.ArgumentParser(description="Add auto eval to training scripts")
    parser.add_argument("--dry-run", action="store_true", help="Only show what would be done")
    parser.add_argument("--base-dir", type=str, default="chengyue_script/dvlm/training",
                        help="Base directory to search for scripts")
    args = parser.parse_args()
    
    base_dir = Path(args.base_dir)
    if not base_dir.exists():
        print(f"Error: Directory not found: {base_dir}")
        return 1
    
    scripts = find_training_scripts(base_dir)
    print(f"Found {len(scripts)} training scripts")
    print()
    
    modified = 0
    skipped_has_eval = 0
    skipped_no_match = 0
    
    for script in scripts:
        content = script.read_text()
        
        # Skip if already has auto eval
        if has_auto_eval(content):
            print(f"⏭️  {script} (already has auto eval)")
            skipped_has_eval += 1
            continue
        
        # Try to add auto eval
        new_content, success = add_auto_eval(content)
        
        if not success:
            print(f"⚠️  {script} (no matching pattern found)")
            skipped_no_match += 1
            continue
        
        if args.dry_run:
            print(f"📝 Would modify: {script}")
        else:
            script.write_text(new_content)
            print(f"✅ Modified: {script}")
        
        modified += 1
    
    print()
    print("=" * 50)
    print(f"Summary:")
    print(f"  Modified: {modified}")
    print(f"  Skipped (already has eval): {skipped_has_eval}")
    print(f"  Skipped (no match): {skipped_no_match}")
    
    if args.dry_run:
        print()
        print("This was a dry run. Use without --dry-run to apply changes.")
    
    return 0


if __name__ == "__main__":
    exit(main())

