#!/usr/bin/env python3
"""
Prompt: 把所有脚本中写死的集群配置改为引用 scripts/cluster_env.sh

Migration script to replace hardcoded cluster configs in shell scripts
with references to scripts/cluster_env.sh.

Replacements:
1. `source ~/.bashrc` + `conda activate XXX || true`
   → `source "${PROJECT_ROOT}/scripts/cluster_env.sh"` + `lmflow_activate_conda`

2. VILA_DATASETS auto-detect block (multi-line if/else)
   → removed (handled by cluster_env.sh)

3. `export VILA_DATASETS="xxx"` (hardcoded)
   → removed (handled by cluster_env.sh, VILA_DATASETS already exported)

4. `export DATASET_BACKEND="webdataset_multi_modal"`
   → removed (handled by cluster_env.sh)

Does NOT touch:
- #SBATCH directives (handled by scripts/submit.sh at submission time)
- third_party/ scripts
- PROJECT_ROOT definitions
"""

import re
import sys
import os
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent

# Directories to migrate (relative to PROJECT_ROOT)
DIRS = [
    "chengyue_script",
    "shiyi_scripts",
]

# Skip third_party and non-.sh files
SKIP_PATTERNS = ["third_party/", "node_modules/", ".git/"]

DRY_RUN = "--dry-run" in sys.argv
VERBOSE = "--verbose" in sys.argv or DRY_RUN


def find_scripts():
    """Find all .sh files in target dirs."""
    scripts = []
    for d in DIRS:
        dirpath = PROJECT_ROOT / d
        if dirpath.is_dir():
            for f in sorted(dirpath.rglob("*.sh")):
                rel = str(f.relative_to(PROJECT_ROOT))
                if not any(skip in rel for skip in SKIP_PATTERNS):
                    scripts.append(f)
    return scripts


def migrate_script(path: Path) -> tuple[bool, list[str]]:
    """
    Migrate a single script. Returns (changed, list_of_changes).
    """
    content = path.read_text()
    original = content
    changes = []

    # --- Pattern 1: Replace top-level `source ~/.bashrc` + `conda activate XXX` ---
    # Multinode scripts: "source ~/.bashrc\nconda activate XXX || true"
    # Some have "source "${HOME}/.bashrc""
    # Only match top-level (indent 0-4 spaces), NOT inside srun bash -c blocks (8+ spaces)
    pattern_bashrc_conda = re.compile(
        r'^([ \t]{0,4})(?:source\s+~/.bashrc|source\s+"\$\{HOME\}/.bashrc")\s*\n'
        r'\1conda\s+activate\s+\S+.*$',
        re.MULTILINE
    )
    if pattern_bashrc_conda.search(content):
        def replace_bashrc_conda(m):
            indent = m.group(1)
            return (
                f'{indent}source "${{PROJECT_ROOT}}/scripts/cluster_env.sh"\n'
                f'{indent}lmflow_activate_conda'
            )
        content = pattern_bashrc_conda.sub(replace_bashrc_conda, content)
        changes.append("replaced source ~/.bashrc + conda activate → cluster_env.sh + lmflow_activate_conda")

    # --- Pattern 2: Remove VILA_DATASETS auto-detect block ---
    # The multi-line if/else block that detects hostname
    pattern_vila_detect = re.compile(
        r'^# VILA dataset registry - auto-detect cluster\n'
        r'# .*\n'
        r'# .*\n'
        r'if \[\[ -z "\$\{VILA_DATASETS:-\}" \]\]; then\n'
        r'.*\n'  # _hostname=...
        r'.*\n'  # if hostname check
        r'.*\n'  # export VILA_DATASETS="oci-nrt-cs"
        r'.*\n'  # echo
        r'.*\n'  # else
        r'.*\n'  # export VILA_DATASETS="cs-oci-ord"
        r'.*\n'  # echo
        r'.*\n'  # fi
        r'else\n'
        r'.*\n'  # echo
        r'fi\n',
        re.MULTILINE
    )
    if pattern_vila_detect.search(content):
        content = pattern_vila_detect.sub(
            '# VILA dataset registry — loaded from cluster_env.sh\n',
            content
        )
        changes.append("removed VILA_DATASETS auto-detect block (now in cluster_env.sh)")

    # --- Pattern 3: Remove hardcoded `export VILA_DATASETS="xxx"` ---
    # (only standalone lines, not inside if/else — those are handled above)
    pattern_vila_export = re.compile(
        r'^([ \t]*)export VILA_DATASETS="[^"]*"\s*$',
        re.MULTILINE
    )
    # But don't remove if it's inside srun's bash -c block (those have deep indentation)
    # and keep lines that use ${...} (they're already variable)
    for m in pattern_vila_export.finditer(content):
        line = m.group(0).strip()
        if "${" in line:
            continue  # already variable, skip
        # Check if this is a top-level export (indent <= 4 spaces)
        indent = m.group(1)
        if len(indent.replace('\t', '    ')) <= 4:
            # Replace with comment
            content = content.replace(
                m.group(0),
                f'{indent}# VILA_DATASETS — loaded from cluster_env.sh'
            )
            changes.append(f"replaced hardcoded VILA_DATASETS export: {line}")

    # --- Pattern 4: Remove `export DATASET_BACKEND="webdataset_multi_modal"` ---
    pattern_backend = re.compile(
        r'^([ \t]*)export DATASET_BACKEND="webdataset_multi_modal"\s*$',
        re.MULTILINE
    )
    for m in pattern_backend.finditer(content):
        indent = m.group(1)
        if len(indent.replace('\t', '    ')) <= 4:
            content = content.replace(
                m.group(0),
                f'{indent}# DATASET_BACKEND — loaded from cluster_env.sh'
            )
            changes.append("replaced hardcoded DATASET_BACKEND export")

    # --- Pattern 5: Handle srun bash -c blocks with source ~/.bashrc + conda activate ---
    # These are inside single-quoted strings, e.g.:
    #   bash -c '
    #       cd ...
    #       source ~/.bashrc
    #       conda activate dvlm || true
    #       ...
    #       export VILA_DATASETS="oci-nrt-cs"
    #       export DATASET_BACKEND="webdataset_multi_modal"
    # Replace with:  source "${PROJECT_ROOT}/scripts/cluster_env.sh" + lmflow_activate_conda
    pattern_srun_bashrc = re.compile(
        r"^([ \t]+)source ~/\.bashrc\n"
        r"\1conda activate \S+ \|\| true",
        re.MULTILINE
    )
    # Only replace matches that occur AFTER first 'srun' or 'bash -c' (i.e. inside srun blocks)
    # We already replaced the top-level ones in Pattern 1, so now find remaining ones
    remaining_matches = list(pattern_srun_bashrc.finditer(content))
    if remaining_matches:
        for m in reversed(remaining_matches):  # reverse to preserve offsets
            indent = m.group(1)
            # Inside srun bash -c '...', PROJECT_ROOT is injected via '"${VAR}"' pattern
            # e.g.  source '"${PROJECT_ROOT}"'/scripts/cluster_env.sh
            replacement = (
                indent + "source '\"${PROJECT_ROOT}\"'/scripts/cluster_env.sh\n" +
                indent + "lmflow_activate_conda"
            )
            content = content[:m.start()] + replacement + content[m.end():]
        changes.append(f"replaced {len(remaining_matches)} srun-block source ~/.bashrc + conda activate")

    # Replace hardcoded VILA_DATASETS inside srun blocks (deeply indented, > 4 spaces)
    pattern_srun_vila = re.compile(
        r'^([ \t]{8,})export VILA_DATASETS="[^"]*"\s*$',
        re.MULTILINE
    )
    srun_vila_matches = list(pattern_srun_vila.finditer(content))
    if srun_vila_matches:
        for m in reversed(srun_vila_matches):
            indent = m.group(1)
            content = content[:m.start()] + f'{indent}# VILA_DATASETS — from cluster_env.sh' + content[m.end():]
        changes.append(f"replaced {len(srun_vila_matches)} srun-block VILA_DATASETS exports")

    # Replace hardcoded DATASET_BACKEND inside srun blocks
    pattern_srun_backend = re.compile(
        r'^([ \t]{8,})export DATASET_BACKEND="webdataset_multi_modal"\s*$',
        re.MULTILINE
    )
    srun_backend_matches = list(pattern_srun_backend.finditer(content))
    if srun_backend_matches:
        for m in reversed(srun_backend_matches):
            indent = m.group(1)
            content = content[:m.start()] + f'{indent}# DATASET_BACKEND — from cluster_env.sh' + content[m.end():]
        changes.append(f"replaced {len(srun_backend_matches)} srun-block DATASET_BACKEND exports")

    # --- Pattern 6: Handle evaluation scripts that use EVAL_CONDA_ENV pattern ---
    # These already source ~/.bashrc conditionally; replace with cluster_env.sh
    pattern_eval_conda = re.compile(
        r'EVAL_CONDA_ENV=\$\{EVAL_CONDA_ENV:-\w+\}\n'
        r'EVAL_CONDA_ACTIVATE=\$\{EVAL_CONDA_ACTIVATE:-1\}\n'
        r'if \[ "\$\{EVAL_CONDA_ACTIVATE\}" != "0" \]; then\n'
        r'    if \[ -f "\$\{HOME\}/.bashrc" \]; then\n'
        r'        source "\$\{HOME\}/.bashrc"\n'
        r'    fi\n'
        r'    if command -v conda >/dev/null 2>&1; then\n'
        r'        conda activate "\$\{EVAL_CONDA_ENV\}" \|\| true\n'
        r'    fi\n'
        r'fi',
        re.MULTILINE
    )
    if pattern_eval_conda.search(content):
        content = pattern_eval_conda.sub(
            '# Conda env — loaded from cluster_env.sh\n'
            'source "${PROJECT_ROOT}/scripts/cluster_env.sh"\n'
            'lmflow_activate_conda',
            content
        )
        changes.append("replaced eval conda activation block → cluster_env.sh")

    changed = content != original
    if changed and not DRY_RUN:
        path.write_text(content)

    return changed, changes


def main():
    scripts = find_scripts()
    print(f"Found {len(scripts)} scripts to check")
    if DRY_RUN:
        print("** DRY RUN — no files will be modified **\n")

    migrated = 0
    for script in scripts:
        changed, changes = migrate_script(script)
        if changed:
            migrated += 1
            rel = script.relative_to(PROJECT_ROOT)
            print(f"{'[DRY] ' if DRY_RUN else ''}Migrated: {rel}")
            for c in changes:
                print(f"  - {c}")

    print(f"\n{'Would migrate' if DRY_RUN else 'Migrated'} {migrated}/{len(scripts)} scripts")


if __name__ == "__main__":
    main()
