from ops import *
import ast, signal, sys, yaml

label=sys.argv[1]; assert label in ('baseline','source-baseline','source-baseline-r2','candidate')
mode='source-baseline' if label.startswith('source-baseline') else label
BASE=HERE; HERE=BASE/label; HERE.mkdir(mode=0o700)
def save(name,data,private=False):
    with ((PRIVATE if private else HERE)/name).open('x') as f: json.dump(data,f,indent=2)
ORDER=[3,5,6,8,9,10,11,12]
LEVELS={'baseline':[6,8,6], 'source-baseline':[8], 'candidate':[8,6]}[mode]
DURATIONS={'baseline':[60,180,60], 'source-baseline':[45], 'candidate':[120,30]}[mode]
PROBE='/opt/contextra/tools/provider_scale_probe.py'
original=yaml.safe_load((PRIVATE/'runtime-original.yaml').read_text())['added']
specs={c:next(s for s in original if str(s.get('camera_id'))==str(c)) for c in ORDER}
sids={c:specs[c]['stream_id'] for c in ORDER};reverse={v:k for k,v in sids.items()}
FIELDS=('intrusion','loitering','fall','fight','fire','crowd','tracking')
saved={c:api('/api/v1/streams/'+sids[c]+'/settings')['configured'] for c in ORDER}
save('settings-'+label+'.json',saved,private=True)
starts={};touched=[];launched=[];active=[];begin=time.monotonic();initial_oom=None;control_pid=None

# Reuse read-only probes and existing safety/settings operations, not the old trial body.
tree=ast.parse(Path('/tmp/codex/ltb1021/2026-09-26_provider-scale_r01/trial.py').read_text())
names={'event','interrupt','patch_settings','basic_safety','sample','model_config','frame_probe'}
exec(compile(ast.Module(body=[n for n in tree.body if isinstance(n,ast.FunctionDef) and n.name in names],type_ignores=[]),'<reused-probes>','exec'))

def rows():
    out={}
    for r in api('/api/v1/streams')['streams']:
        c=reverse.get(r['stream_id'])
        if c not in launched: continue
        out[str(c)]={k:r.get(k) for k in ('state','run_id','run_generation','last_error','cleanup_error','cleanup_pending','degraded','scenarios','publisher_metrics')}
        out[str(c)]['providers']={p:{k:v.get(k) for k in ('state','child_pid','last_error','cleanup_error','last_cleanup','metadata_path','publisher_metrics')}
            for p,v in (r.get('provider_processes') or {'pose':r}).items()}
    return out

signal.signal(signal.SIGTERM,interrupt);signal.signal(signal.SIGINT,interrupt)
assert not inspect('contextra-ai-gateway-1')['State']['Running']
result=dict(stages=[],starts=starts,errors=[],cleanup=None)
def persist_inputs(current,where):
    with (HERE/'startup.jsonl').open('a') as f:
        f.write(json.dumps(dict(t=time.monotonic(),wall_ms=time.time_ns()//1000000,where=where,channels=current))+'\n')
try:
    configs=model_config();save('model-configs.json',configs)
    conditions=[]
    for c in ORDER:
        desired={k:dict(enabled=True) for k in ('intrusion','loitering','fall','fight','tracking')}
        desired.update(fire=dict(enabled=False,scenario=dict(enabled=False)),crowd=dict(enabled=False))
        touched.append(c);v=patch_settings(c,desired)['configured']
        assert [p['key'] for p in v['execution_plan']['providers']]==['pose']
        assert v['analysisFpsTarget']==10
        conditions.append(dict(camera=c,source_sha256=hashlib.sha256(specs[c]['uri'].encode()).hexdigest(),
            source_options_sha256=hashlib.sha256(json.dumps(specs[c].get('source_options',{}),sort_keys=True).encode()).hexdigest(),
            original_revision=saved[c]['revision'],trial_revision=v['revision'],scenarios={k:v[k] for k in FIELDS},pose=v['pose']))
    save('conditions.json',dict(levels=LEVELS,durations=DURATIONS,conditions=conditions))
    for index,(level,seconds) in enumerate(zip(LEVELS,DURATIONS)):
        stage=dict(level=level,index=index,seconds=seconds,samples=[],status='RUNNING')
        result['stages'].append(stage);event('stage_begin',level=level,index=index)
        for c in list(active):
            if c in ORDER[:level]: continue
            request('/api/v1/streams/'+sids[c]+'/stop',{},timeout=20)
            end=time.monotonic()+90
            while rows()[str(c)]['cleanup_pending'] or rows()[str(c)]['providers']['pose'].get('child_pid'):
                if time.monotonic()>end:raise RuntimeError('STEPDOWN_REAP_BOUND')
                time.sleep(1)
            active.remove(c)
        for c in ORDER[:level]:
            if c in active:continue
            basic_safety(rows());launched.append(c);active.append(c)
            sent=time.monotonic();wall=time.time_ns()//1000000
            ack=request('/api/v1/streams/'+sids[c]+'/start',{},timeout=20)
            starts[c]=dict(sent=sent,wall_ms=wall,ack_at=time.monotonic(),request_state=ack.get('state'))
            event('start_ack',camera=c,**starts[c]);deadline=time.monotonic()+70
            while True:
                current=rows();basic_safety(current);persist_inputs(current,'startup-'+str(c))
                r=current[str(c)]
                if len(r['providers'])==1 and r['providers']['pose'].get('metadata_path',{}).get('providers',{}).get('pose',{}).get('responses',0)>=2:
                    starts[c]['ready_observed_at']=time.monotonic();break
                if time.monotonic()>deadline:raise RuntimeError('NO_PROVIDER_RESPONSE_'+str(c))
                time.sleep(1)
        time.sleep(5);basic_safety(rows());start=time.monotonic()
        stage['started_wall_ms']=time.time_ns()//1000000
        stage['started_mono']=start
        for i in range(seconds//5+1):
            time.sleep(max(0,start+5*i-time.monotonic()))
            value=sample();basic_safety(value['channels'])
            value['channels']={c:r for c,r in value['channels'].items() if int(c) in active}
            stage['samples'].append(value)
            with (HERE/'samples.jsonl').open('a') as f:f.write(json.dumps(dict(stage=index,level=level,**value))+'\n')
            for r in value['channels'].values():
                path=r['providers']['pose']['metadata_path']
                assert path['execution_path']=='provider-linear-trial' and not path['freshness_enabled'] and path['queue_max_age_s']==1
                timing=path['performance']['pipeline_diagnostic']['providers']['pose']['input_timeline']
                if mode=='candidate':
                    assert path['source_latest_trial'] and timing and any(f['input'].get('latest_applied')==1 for f in timing), 'SOURCE_POLICY_NOT_APPLIED'
            event('sample',level=level,index=i)
        stage['frames']=frame_probe(ORDER[:level]);assert configs==model_config()
        stage['status']='PASS';stage['finished_wall_ms']=time.time_ns()//1000000
        save('stage-'+str(index)+'-'+str(level)+'.json',stage)
        print(json.dumps(dict(level=level,status='PASS',seconds=seconds)),flush=True)
except Exception as e:
    result['errors'].append(dict(type=type(e).__name__,reason=str(e)[:160]))
    event('trial_error',**result['errors'][-1])
finally:
    for c in launched:
        try:request('/api/v1/streams/'+sids[c]+'/stop',{},timeout=20)
        except Exception as e:result['errors'].append(dict(type=type(e).__name__,reason='STOP_REQUEST_'+str(c)))
    deadline=time.monotonic()+90
    while True:
        r=rows();reaped=all(not v['cleanup_pending'] and all(not p.get('child_pid') for p in v['providers'].values()) for v in r.values())
        if reaped or time.monotonic()>deadline:break
        time.sleep(1)
    result['cleanup']=dict(reaped=reaped,states=r)
    if reaped:
        for c in touched:
            try:
                if c in launched and r[str(c)]['state']!='stopped':raise RuntimeError('RESTORE_DEFERRED_ERROR_STATE')
                changes={k:saved[c][k] for k in FIELDS}
                if 'scenario' not in changes['fire']:changes['fire']=dict(changes['fire'],scenario={'enabled':False})
                patch_settings(c,changes)
            except Exception as e:result['errors'].append(dict(type=type(e).__name__,reason='RESTORE_SETTINGS_'+str(c)))
    save('trial-result.json',result)
    print(json.dumps(dict(errors=result['errors'],reaped=reaped)),flush=True)
if result['errors'] or not reaped:raise SystemExit(1)