"""Exclusive-output fresh-process runner; original cases, one science GPU job."""
import argparse
import json
import os
import urllib.request
from context import ROOT, MODEL, EXP, AG, AG_MODEL, CFG, check, budget, read, load_module


def preflight():
    os.environ['CUDA_VISIBLE_DEVICES'] = CFG['gpu_uuid']
    os.environ['CUDA_HOME'] = '/usr/local/cuda'
    os.environ['PATH'] = str(ROOT / 'models/male-cns-lif/.venv/bin') + ':/usr/local/cuda/bin:' + os.environ['PATH']
    os.environ['MAX_JOBS'] = '2'
    os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'
    for port in (8188, 8189, 8190):
        with urllib.request.urlopen(f'http://192.168.10.51:{port}/queue', timeout=5) as response:
            queue = json.load(response)
        assert not queue['queue_running'] and not queue['queue_pending'], 'ComfyUI busy: defer, never evict'


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('backend', choices=['gpu', 'oracle', 'component'])
    parser.add_argument('arm', choices=['delay0', 'delay18'])
    parser.add_argument('case')
    parser.add_argument('--repeat', type=int, default=0)
    args = parser.parse_args()
    assert '/' not in args.case and '..' not in args.case and 0 <= args.repeat < 3
    check(); budget()
    assert read(EXP / 'origin-audit/summary.json')['pass_all']
    assert read(EXP / 'audit/shift-gate.json')['pass_all']
    preflight()
    case = EXP / 'cases' / args.arm / args.case
    assert case.is_dir() and read(case / 'spec.json')['apl_delay_steps'] == int(args.arm[5:])
    if args.case not in CFG['fixtures']:
        assert read(EXP / 'audit/fixtures-gate.json')['pass_all'], 'Fixture gate required before fullgraph execution'
        if args.case not in CFG['calibration']['full_oracle_probes'] or args.repeat > 0:
            assert read(EXP / 'audit/probes-gate.json')['pass_all'], 'Both fullgraph probe gates required before CAL'
    out = EXP / 'gpu' / args.arm / f'{args.case}-r{args.repeat}'
    if args.backend == 'gpu':
        from gpu_timing import run_gpu
        result = run_gpu(case, out, EXP / 'extension-build')
    elif args.backend == 'oracle':
        assert args.repeat == 0
        from oracle_timing import run_case
        result = run_case(case, EXP / 'oracle' / args.arm / args.case)
    else:
        assert args.repeat == 0
        component = load_module('sealed_ag_component_aj', AG_MODEL / 'validate_components.py')
        result = component.validate(out, AG / 'apl-bundle', EXP / 'component-audit' / args.arm / (args.case + '.json'))
    budget()
    print(json.dumps(dict(backend=args.backend, arm=args.arm, case=args.case, repeat=args.repeat,
        complete=result.get('complete'), pass_all=result.get('pass_all')), allow_nan=False), flush=True)
    if not result.get('pass_all', False):
        raise SystemExit(1)


if __name__ == '__main__':
    main()