from ops import *
import sys
phase=sys.argv[1]
targets=json.loads((HERE/'targets.json').read_text())
if phase.endswith('single'):targets=targets[:1]
count=len(targets)
assert phase in ('old-single','old-ten','current-single','current-ten')
def triton():
    code='''import urllib.request,json
models=['RFDETRMedium','RFDETRKeypointPreview','RFDETRFireSmoke','Awiros_person_and_head_detection']
print(json.dumps({m:dict(stats=json.load(urllib.request.urlopen('http://localhost:8000/v2/models/'+m+'/stats',timeout=3))['model_stats'][0],config=json.load(urllib.request.urlopen('http://localhost:8000/v2/models/'+m+'/config',timeout=3))) for m in models}))'''
    return json.loads(cmd(['docker','exec','contextra-analytics-1','python3','-B','-c',code],timeout=15))
def cpu():
    names=cmd(['docker','top','contextra-analytics-1','-eo','pid,comm']).splitlines()[1:]
    mapping={r['child_pid']:str(next(t['camera'] for t in targets if t['stream_id']==r['stream_id'])) for r in api('/api/v1/streams')['streams'] if r['state']=='running'}
    out={}
    for row in names:
        pid,name=row.split()
        if name not in ('python3','tritonserver'):continue
        root=Path('/proc')/pid
        try:ns=int(next(x for x in (root/'status').read_text().splitlines() if x.startswith('NSpid:')).split()[-1])
        except (OSError,StopIteration):continue
        label='Triton' if name=='tritonserver' else mapping.get(ns,'parent')
        for t in (root/'task').glob('*'):
            try:out[pid+':'+t.name]=dict(channel=label,values=[int(v) for v in (t/'schedstat').read_text().split()])
            except OSError:pass
    return out
def sample():
    rows=api('/api/v1/streams')['streams']
    return dict(t=time.monotonic(),channels=[dict(camera=t['camera'],**{k:r.get(k) for k in ('state','run_id','child_pid','last_error','metadata_path','publisher_metrics')}) for t in targets for r in rows if t['stream_id']==r['stream_id']],
        publisher=api('/api/v1/metadata/status')['metrics'],
        cpu=Path('/proc/stat').read_text().splitlines()[0],pressure=Path('/proc/pressure/cpu').read_text(),
        gpu=cmd(['nvidia-smi','--query-gpu=utilization.gpu,memory.used','--format=csv,noheader,nounits']).strip())
assert not inspect('contextra-ai-gateway-1')['State']['Running']
save(phase+'-intent.json',dict(cameras=[t['camera'] for t in targets],seconds=30,image=inspect('contextra-analytics-1')['Image']))
out={};samples=[];started=[]
try:
    for t in targets:
        sid=t['stream_id'];started.append(sid);request('/api/v1/streams/'+sid+'/start',{},timeout=15)
    deadline=time.monotonic()+60;stable=None
    while True:
        rows=[api('/api/v1/streams/'+t['stream_id']) for t in targets]
        if any(r['state']=='error' for r in rows):raise RuntimeError('CHANNEL_ERROR')
        ready=all(r['state']=='running' and r.get('metadata_path',{}).get('performance',{}).get('pipeline_diagnostic',{}).get('counts',{}).get('copied',0)>=5 for r in rows)
        if ready:
            stable=stable or time.monotonic()
            if time.monotonic()-stable>=5:break
        if time.monotonic()>deadline:raise RuntimeError('WARMUP_FAILED')
        time.sleep(1)
    tri0=triton();c0=cpu();start=time.monotonic()
    for i in range(7):
        time.sleep(max(0,start+5*i-time.monotonic()))
        s=sample();samples.append(s)
        if any(r['state']!='running' for r in s['channels']):raise RuntimeError('CHANNEL_LOST')
    c1=cpu();tri1=triton();out.update(triton_before=tri0,triton_after=tri1,cpu_before=c0,cpu_after=c1,seconds=samples[-1]['t']-samples[0]['t'],samples=samples)
    assert all(a['run_id']==b['run_id'] for a,b in zip(samples[0]['channels'],samples[-1]['channels']))
except Exception as e:
    out['error']=type(e).__name__+':'+str(e)[:100]
finally:
    for sid in started:request('/api/v1/streams/'+sid+'/stop',{},timeout=90)
    deadline=time.monotonic()+25
    while True:
        rows=api('/api/v1/streams')['streams']
        if all(r['state']=='stopped' and not r.get('child_pid') and not r.get('cleanup_pending') for r in rows):break
        if time.monotonic()>deadline:raise RuntimeError('STOP_REAP_FAILED')
        time.sleep(1)
    out['stopped']=len(rows);save(phase+'.json',out)
    print(json.dumps(dict(phase=phase,seconds=out.get('seconds'),error=out.get('error'),stopped=len(rows))),flush=True)
if out.get('error'):raise SystemExit(1)