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

import os
import subprocess
import sys
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(argv: list[str]) -> int:
    if len(argv) != 2:
        print('usage: run_quick_probe_on_cluster.py CLUSTER', file=sys.stderr)
        return 2
    cluster = argv[1]
    env = os.environ.copy()
    env.pop('PYTHONPATH', None)
    env['HF_TOKEN'] = _secret('HF_TOKEN')
    env['WANDB_API_KEY'] = _secret('WANDB_API_KEY')
    remote_cmd = (
        'uv run --no-dev python scripts/aggressive_oom_node_wrapper.py '
        '--jobs artifacts/aggressive_oom_sweep_quick_probe/jobs.jsonl '
        '--node quick_probe_node '
        '--artifact-dir artifacts/current_run_quick_probe '
        '--utilization artifacts/aggressive_oom_sweep_quick_probe/utilization.jsonl '
        '--max-jobs 1 '
        '--sample-interval-seconds 5 '
        '--require-success'
    )
    command = [
        'uv',
        'run',
        'sky',
        'exec',
        '--workdir',
        '.',
        '--gpus',
        'RTX4090:1',
        '--secret',
        'HF_TOKEN',
        '--secret',
        'WANDB_API_KEY',
        cluster,
        remote_cmd,
    ]
    print(f'quick_probe_cluster={cluster}', flush=True)
    print('quick_probe_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(sys.argv))