"""One registered AK action per fresh process; no changes to frozen runtimes."""
import argparse
import json
import os
import urllib.request

from context import ROOT, EXP, CFG, check, budget, read
import runtime_bridge as bridge


def preflight():
    # Establish visibility before the lazy runtime/component Torch import.
    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)
        bridge.require(not queue['queue_running'] and not queue['queue_pending'], 'ComfyUI busy: defer, never evict')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('backend', choices=['gpu', 'oracle', 'component'])
    parser.add_argument('group', choices=bridge.GROUPS)
    parser.add_argument('case')
    parser.add_argument('--repeat', type=int, default=0)
    args = parser.parse_args()
    check(); budget()
    row = bridge.authorize_case(EXP, CFG, args.group, args.case, args.backend, args.repeat)
    case = EXP / 'cases' / args.group / args.case
    spec = bridge.case_spec(case)
    bridge.require(spec['name'] == args.case and spec['budget_root'] == str(EXP) and
                   spec['output_budget_bytes'] == CFG['output_budget_bytes'], 'Case identity/new budget differs')
    if args.group == 'controls':
        bridge.control_inputs(case, row['source_case'])
    else:
        bridge.require(spec.get('kind') == 'fullgraph' and spec.get('validation_cohort') == args.group and
                       all(spec[key] == row[key] for key in ('mode', 'odor_id', 'odor_index', 'seed')),
                       'New fullgraph case differs from its registered cohort cell')
    output = EXP / 'gpu' / args.group / f'{args.case}-r{args.repeat}'
    if args.backend == 'component':
        summary = read(output / 'summary.json')
        bridge.require(summary.get('complete') is True and summary.get('pass_all') is True and
                       summary.get('case') == args.case and summary.get('apl_delay_steps') == 18,
                       'Component audit needs a complete matching D18 GPU run')
    preflight()
    def execute():
        if args.backend == 'gpu':
            result = bridge.run_gpu(case, output, EXP / 'extension-build')
        elif args.backend == 'oracle':
            result = bridge.run_oracle(case, EXP / 'oracle' / args.group / args.case)
        else:
            result = bridge.component(output, spec['apl_bundle_dir'], EXP / 'component-audit' / args.group / (args.case + '.json'))
        budget()
        bridge.require(result.get('pass_all') is True, 'Runtime/component evidence did not pass its own integrity checks')
        return result
    result = bridge.guarded_action(EXP, args.backend, args.group, args.case, args.repeat, execute)
    print(json.dumps(dict(backend=args.backend, group=args.group, case=args.case, repeat=args.repeat,
        complete=result.get('complete'), pass_all=result.get('pass_all')), allow_nan=False), flush=True)


if __name__ == '__main__':
    main()