#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
import subprocess
from pathlib import Path


def _secret(name: str) -> str:
    value = os.environ.get(name)
    if value:
        return value
    path = Path('.env') / name
    if path.is_file():
        value = path.read_text().strip()
        if value:
            return value
    raise RuntimeError(f'{name} is required')


def main() -> int:
    parser = argparse.ArgumentParser(description='Run an aggressive sweep node on an existing Sky cluster.')
    parser.add_argument('cluster')
    parser.add_argument('--node', required=True)
    parser.add_argument('--jobs', default='artifacts/aggressive_oom_sweep/jobs.jsonl')
    parser.add_argument('--artifact-dir', default='artifacts/current_run')
    parser.add_argument('--utilization', default='artifacts/aggressive_oom_sweep/utilization.jsonl')
    parser.add_argument('--max-jobs', type=int)
    parser.add_argument('--sample-interval-seconds', type=float, default=30.0)
    parser.add_argument('--stale-after-seconds', type=float, default=21600.0)
    parser.add_argument('--max-attempts', type=int, default=2)
    args = parser.parse_args()

    env = os.environ.copy()
    env.pop('PYTHONPATH', None)
    env['HF_TOKEN'] = _secret('HF_TOKEN')
    env['WANDB_API_KEY'] = _secret('WANDB_API_KEY')
    parts = [
        'uv run --no-dev python scripts/aggressive_oom_node_wrapper.py',
        f'--jobs {args.jobs}',
        f'--node {args.node}',
        f'--artifact-dir {args.artifact_dir}',
        f'--utilization {args.utilization}',
        f'--sample-interval-seconds {args.sample_interval_seconds:g}',
        f'--stale-after-seconds {args.stale_after_seconds:g}',
        f'--max-attempts {args.max_attempts}',
    ]
    if args.max_jobs is not None:
        parts.append(f'--max-jobs {args.max_jobs}')
    remote_cmd = ' '.join(parts)
    command = [
        'uv',
        'run',
        'sky',
        'exec',
        '--workdir',
        '.',
        '--gpus',
        'RTX4090:1',
        '--secret',
        'HF_TOKEN',
        '--secret',
        'WANDB_API_KEY',
        args.cluster,
        remote_cmd,
    ]
    print(f'sweep_cluster={args.cluster}', flush=True)
    print(f'sweep_node={args.node}', flush=True)
    print('sweep_command=sky exec --workdir . --gpus RTX4090:1 --secret HF_TOKEN --secret WANDB_API_KEY <cluster> <wrapper>', flush=True)
    return subprocess.run(command, env=env, check=False).returncode


if __name__ == '__main__':
    raise SystemExit(main())