"""Supervise the one-model job; no automatic download of subsequent models."""
import argparse
import json
import os
import pathlib
import subprocess
import sys
import time


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--root', type=pathlib.Path, required=True)
    args = p.parse_args()
    root = args.root.resolve()
    root.mkdir(exist_ok=True, parents=True)
    status_file = root / 'status.json'
    start = time.time()
    env = dict(os.environ, OPENBLAS_NUM_THREADS='1', OMP_NUM_THREADS='1')

    def status(phase, **extra):
        status_file.write_text(json.dumps({'phase': phase, 'updated_unix': time.time(),
                                           'supervisor_started_unix': start, **extra}, indent=2))
        print(phase, flush=True)

    def command(module, log, timeout, extra=()):
        with (root / log).open('w') as stream:
            subprocess.run([sys.executable, '-B', '-m', module, '--root', str(root), *extra],
                           check=True, timeout=timeout, env=env, stdout=stream, stderr=subprocess.STDOUT)

    try:
        status('waiting_for_download')
        download_wait_start = time.monotonic()
        while True:
            manifest_path = root / 'manifest.json'
            if manifest_path.exists():
                manifest = json.loads(manifest_path.read_text())
                if manifest.get('completed_unix') and all((root / 'weights' / f['name']).exists() for f in manifest['files']):
                    break
            if time.monotonic() - download_wait_start > 7200:
                raise TimeoutError('Download wait exceeded 2 hours; no analysis started')
            time.sleep(10)
        status('analyzing', analysis_hard_limit_seconds=7200)
        command('pilot.analyze', 'analysis.log', 7200, ('--seconds', '6600'))
        status('plotting')
        command('pilot.plot', 'plot.log', 300)
        status('validating')
        command('pilot.validate', 'validate.log', 120)
        with (root / 'environment.txt').open('w') as f:
            subprocess.run([sys.executable, '-B', '-m', 'pip', 'freeze'], check=True, stdout=f, env=env)
        status('results_ready_weights_retained_until_archive_copy')
    except Exception as e:
        status('failed', error_type=type(e).__name__, error=str(e))
        raise


if __name__ == '__main__':
    main()