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

import argparse
import signal
import subprocess
import time
from pathlib import Path

_stop = False


def _handle_stop(signum: int, frame: object) -> None:
    global _stop
    _stop = True


def _run(command: list[str]) -> int:
    completed = subprocess.run(command, check=False)
    return int(completed.returncode)


def _rsync(host: str, port: int, remote: str, local: Path, *, include_metrics_only: bool = False) -> int:
    local.mkdir(parents=True, exist_ok=True)
    command = [
        'rsync',
        '-az',
        '--delete',
        '--exclude',
        '*.pt',
        '--exclude',
        '.wandb/',
    ]
    if include_metrics_only:
        command.extend([
            '--include',
            '*/',
            '--include',
            '*.json',
            '--include',
            '*.jsonl',
            '--include',
            '*.toml',
            '--include',
            '*.txt',
            '--exclude',
            '*',
        ])
    command.extend([
        '-e',
        f'ssh -p {port} -o StrictHostKeyChecking=no',
        f'root@{host}:{remote}',
        str(local),
    ])
    return _run(command)


def main() -> int:
    parser = argparse.ArgumentParser(description='Continuously rsync lightweight sweep artifacts from one node.')
    parser.add_argument('--host', required=True)
    parser.add_argument('--port', type=int, required=True)
    parser.add_argument('--remote-root', default='/root/sky_workdir')
    parser.add_argument('--local-root', default='artifacts/remote_runs/live_node0')
    parser.add_argument('--interval-seconds', type=float, default=60.0)
    args = parser.parse_args()

    signal.signal(signal.SIGINT, _handle_stop)
    signal.signal(signal.SIGTERM, _handle_stop)
    local_root = Path(args.local_root)
    remote_root = args.remote_root.rstrip('/')
    interval = max(5.0, args.interval_seconds)
    cycle = 0
    while not _stop:
        started = time.time()
        cycle += 1
        rcodes = [
            _rsync(args.host, args.port, f'{remote_root}/artifacts/aggressive_oom_sweep/', local_root / 'aggressive_oom_sweep/'),
            _rsync(args.host, args.port, f'{remote_root}/artifacts/current_run/', local_root / 'current_run/'),
            _rsync(args.host, args.port, f'{remote_root}/artifacts/runs/', local_root / 'runs/', include_metrics_only=True),
        ]
        print(f'collector_cycle={cycle} rcodes={rcodes} elapsed_seconds={time.time() - started:.1f}', flush=True)
        deadline = time.time() + interval
        while not _stop and time.time() < deadline:
            time.sleep(min(1.0, deadline - time.time()))
    print('collector_stopped=true', flush=True)
    return 0


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