from ops import *
COMPOSE=['docker','compose','--env-file',str(ENV),'-f',str(REPO/'docker-compose.yml'),'-f',str(ACTIVE)]
MODELS=['RFDETRMedium','RFDETRKeypointPreview','RFDETRFireSmoke']
def atomic(path,data,expected):
    assert file_hash(path)==expected
    st=path.stat();fd,tmp=tempfile.mkstemp(prefix='.pose-fp16-',dir=path.parent)
    try:
        os.fchmod(fd,st.st_mode&0o777);os.fchown(fd,st.st_uid,st.st_gid)
        with os.fdopen(fd,'wb') as f:f.write(data);f.flush();os.fsync(f.fileno())
        assert file_hash(path)==expected;os.replace(tmp,path)
    finally:
        if os.path.exists(tmp):os.unlink(tmp)
def prepare():
    p=read('plan.json');rt=Runtime(p['config'],R);assert not NEW.exists()
    assert file_hash(ACTIVE)==p['active_hash'] and file_hash(ENV)==p['env_hash']
    assert file_hash(SOURCE)==p['onnx_hash']
    engine=read('engine.json');image=read('image.json')
    code="""from pathlib import Path
from google.protobuf import text_format
from tritonclient.grpc import model_config_pb2 as pb
m=pb.ModelConfig();text_format.Parse(Path('/opt/contextra/models/triton/RFDETRKeypointPreview/config.pbtxt').read_text(),m)
assert m.input[0].data_type==pb.TYPE_FP32
m.version_policy.specific.versions[:]=[4];m.max_batch_size=4;m.input[0].data_type=pb.TYPE_FP16
del m.instance_group[:];g=m.instance_group.add();g.count=8;g.kind=pb.ModelInstanceGroup.KIND_GPU;g.gpus.append(0)
m.dynamic_batching.preferred_batch_size[:]=[4];m.dynamic_batching.max_queue_delay_microseconds=5000
assert all(x.data_type==pb.TYPE_FP32 for x in m.output)
print(text_format.MessageToString(m))
"""
    config=cmd(['docker','exec',rt.container,'python3','-c',code])
    if TARGET.exists():
        assert file_hash(TARGET/'4/model.plan')==engine['sha256']
        assert (TARGET/'config.pbtxt').read_text()==config
    else:
        (TARGET/'4').mkdir(parents=True,mode=0o700);shutil.copy2(R/'pose-fp16.plan',TARGET/'4/model.plan')
        (TARGET/'config.pbtxt').write_text(config)
    for row in p['rows']:
        if row['original'] and not (R/('old-stop-'+row['stream_id']+'.json')).exists():save('old-stop-'+row['stream_id']+'.json',rt.stop(row['stream_id']))
    assert all(q['state']=='stopped' for q in rt.states().values())
    for row in p['rows']:
        sid=row['stream_id'];before=p['settings'][sid];actual=rt.settings(sid)
        assert before['display_categories']==['person']
        body=dict(expected_revision=before['revision'],pose=dict(enabled=True),intrusion=dict(enabled=False))
        marker=R/('patch-'+hashlib.sha256(sid.encode()).hexdigest()[:16]+'.json')
        if marker.exists():
            assert actual['revision']!=before['revision']
            value=dict(configured=actual,restart_pending=False,resumed_readback=True)
        else:
            assert actual==before
            with marker.open('x') as f:json.dump(dict(before=before,request=body),f)
            value=rt.api.call('PATCH','/api/v1/streams/'+quote(sid,safe='')+'/settings',body)
        expected=copy.deepcopy(before);expected['pose']['enabled']=True;expected['intrusion']['enabled']=False
        expected['revision']=value['configured']['revision']
        derived={'execution_plan','categories'}
        assert {k:v for k,v in value['configured'].items() if k not in derived}=={k:v for k,v in expected.items() if k not in derived}
        assert value['configured']['execution_plan']==dict(category='person',model='RFDETRKeypointPreview',role='pose',outputs=['bbox','keypoints'])
        assert not value['restart_pending']
        save('applied-'+hashlib.sha256(sid.encode()).hexdigest()[:16]+'.json',value)
    cfg=copy.deepcopy(p['config']);cfg['expected_image']=image['image']
    cfg['model'].update(name='RFDETRKeypointPreview',version=4,role='pose',instances=8,max_batch=4,queue_delay_us=5000)
    cfg['settings_expected']['pose']['enabled']=True;cfg['settings_expected']['intrusion']['enabled']=False
    cfg['failure_stop_all']=True
    cfg['artifacts']=[dict(name='engine',path=str(TARGET/'4/model.plan'),sha256=engine['sha256']),
        dict(name='model-config',path=str(TARGET/'config.pbtxt'),sha256=file_hash(TARGET/'config.pbtxt')),
        dict(name='registry',path=cfg['registration_file'],sha256=file_hash(cfg['registration_file']))]
    NEW.write_text(yaml.safe_dump(cfg,sort_keys=False));NEW.chmod(0o600)
    save('prepared.json',dict(config=str(NEW),settings={r['stream_id']:rt.settings(r['stream_id']) for r in p['rows']}))
    print('All16 Pose settings applied before Start',flush=True)
def reflect():
    p=read('plan.json');rt=Runtime(load_config(NEW),R);assert not (R/'reflect-start.json').exists()
    save('reflect-start.json',dict(at=time.time(),image=read('image.json')['image']))
    active=yaml.safe_load(ACTIVE.read_text());a=active['services']['analytics'];a['image']=read('image.json')['image']
    for v in a['volumes']:
        if v['target']=='/opt/contextra/models/triton/RFDETRKeypointPreview':v['source']=str(TARGET)
    atomic(ACTIVE,yaml.safe_dump(active,sort_keys=False).encode(),p['active_hash'])
    cmd(COMPOSE+['config','--quiet']);cmd(COMPOSE+['up','--no-start','--no-deps','--force-recreate','--pull','never','analytics'],120)
    cmd(['docker','start',rt.container]);deadline=time.monotonic()+120
    while time.monotonic()<deadline:
        actual=rt.inspect()
        if actual['State'].get('Health',{}).get('Status')=='healthy':break
        time.sleep(1)
    else:raise RuntimeError('Readiness deadline')
    assert actual['Image']==read('image.json')['image']
    assert dict(x.split('=',1) for x in actual['Config']['Env'])==dict(x.split('=',1) for x in p['inspect']['Config']['Env'])
    for k in ('NanoCpus','CpuQuota','CpusetCpus','Memory','MemorySwap','PidMode','RestartPolicy'):
        assert actual['HostConfig'][k]==p['inspect']['HostConfig'][k]
    before={m['Destination']:m['Source'] for m in p['inspect']['Mounts']};after={m['Destination']:m['Source'] for m in actual['Mounts']}
    before['/opt/contextra/models/triton/RFDETRKeypointPreview']=str(TARGET);assert before==after
    for n in ('RFDETRMedium','RFDETRFireSmoke'):assert rt.triton('v2/models/'+n+'/config')==p['model_configs'][n]
    for path,digest in read('image.json')['hashes'].items():
        if path.startswith('app/'):assert cmd(['docker','exec',rt.container,'sha256sum','/opt/contextra/'+path]).split()[0]==digest
    save('reflected.json',dict(image=actual['Image'],mounts=actual['Mounts'],model=rt.triton('v2/models/RFDETRKeypointPreview/config')))
    rt.resources();print('Pose model ready, Medium/Fire configs preserved',flush=True)
def contract():
    code="""import json,numpy as np
import tritonclient.grpc as grpc
c=grpc.InferenceServerClient('localhost:8001');out=[]
for b in range(1,5):
    x=np.random.default_rng(100+b).uniform(-2,2,(b,3,576,576)).astype(np.float16)
    t=grpc.InferInput('input',x.shape,'FP16');t.set_data_from_numpy(x)
    r=c.infer('RFDETRKeypointPreview',[t],model_version='4',request_id='contract-'+str(b),client_timeout=30)
    assert r.get_response().id=='contract-'+str(b);item={'batch':b,'outputs':{}}
    for n,shape in [('dets',(b,100,4)),('labels',(b,100,2)),('keypoints',(b,100,34,8))]:
        y=r.as_numpy(n);assert y.shape==shape and y.dtype==np.float32 and np.isfinite(y).all()
        item['outputs'][n]={'shape':list(y.shape),'dtype':str(y.dtype),'finite':True}
    out.append(item)
print(json.dumps(out))
"""
    save('contract.json',dict(batches=json.loads(cmd(['docker','exec','contextra-analytics-1','python3','-c',code],150)),quality='not evaluated'))
class PoseTrial(Benchmark):
    def metrics(self,rows,enabled):
        states=self.runtime.states();out={}
        for row in rows:
            sid=row['stream_id'];q=states[sid]
            if q['state']!='running':continue
            try:out[sid]=self.runtime.api.call('POST','/api/v1/streams/'+quote(sid,safe='')+'/metrics',dict(run_generation=q['run_generation'],enabled=enabled,seconds=60))
            except Exception as e:out[sid]=dict(error=type(e).__name__)
        self.runtime.save(f'metrics-{len(rows)}-'+('on' if enabled else 'off')+'.json',out)
    def observe(self,rows,seconds):
        self.metrics(rows,True)
        try:
            out=super().observe(rows,seconds);state=self.runtime.states();frames={};checks={}
            for row in rows:
                sid=row['stream_id'];key='contextra:metadata:v1:camera:'+hashlib.sha256(row['source_id'].encode()).hexdigest()+':frames'
                entries=self.runtime.redis('XREVRANGE',key,'+','-','COUNT','2');values=[]
                for e in entries:
                    v=json.loads(dict(zip(e[1][::2],e[1][1::2]))['json'])
                    assert v['sourceId']==row['source_id'] and v['runId']==state[sid]['run_id'] and v['streamId']==sid
                    assert v['coordinateSpace']=='SOURCE_FRAME' and v['normalized']
                    for o in v['objects']:
                        assert o['modelKey']=='RFDETRKeypointPreview' and o['keypointSchema']=='COCO_17'
                        assert len(o['keypoints'])==17 and len(o['keypointInFrame'])==17
                        assert all(len(k)==3 and all(math.isfinite(x) for x in k) for k in o['keypoints'])
                        assert not o['eventIds'] and 'scenarioResults' not in o
                    values.append(v)
                if len(values)==2:assert values[0]['frameSeq']>values[1]['frameSeq']
                sets=[{str(o['trackId']) for o in v['objects'] if o['trackId'] is not None} for v in values]
                frames[sid]=values;checks[sid]=dict(frames=len(values),objects=[len(v['objects']) for v in values],
                    tracked=[len(x) for x in sets],adjacent_track_ids=len(set.intersection(*sets)) if sets else 0,
                    pose_contract=True,quality='not evaluated')
            self.runtime.save(f'metadata-{len(rows)}-private.json',frames);self.runtime.save(f'metadata-{len(rows)}-summary.json',checks)
            w=json.loads((self.runtime.evidence/f'window-{len(rows)}.json').read_text())
            out['bounded_durations_last']={s:v['durations'] for s,v in w['samples'][-1]['channels'].items()}
            return out
        finally:self.metrics(rows,False)
def run():
    cfg=load_config(NEW);e=R/'ascent';e.mkdir(mode=0o700);rt=Runtime(cfg,e)
    with Path(cfg['lock_file']).open('a+') as f:
        fcntl.flock(f,fcntl.LOCK_EX|fcntl.LOCK_NB);p=preflight(rt)
        rt.save('plan.json',p);rt.save('conditions.json',dict(config=cfg,head=read('plan.json')['head'],
            tools={f.name:file_hash(f) for f in (REPO/'tools').glob('stream_benchmark*.py')}))
        save('unused-model-before.json',{n:rt.triton('v2/models/'+n+'/stats') for n in ('RFDETRMedium','RFDETRFireSmoke')})
        save('ascent-start.json',dict(at=time.time(),count=1))
        if PoseTrial(rt).run(p):raise RuntimeError('Ascent failed; do not restore/restart')
def finish(error):
    cfg=load_config(NEW if NEW.exists() else OLD);rt=Runtime(cfg,R/'ascent' if (R/'ascent').exists() else R);p=read('plan.json');cleanup={}
    if error:
        for row in reversed(p['rows']):
            sid=row['stream_id'];q=rt.states()[sid]
            if not q.get('child_pid') and q['state'] in ('stopped','error'):cleanup[sid]=q;continue
            try:cleanup[sid]=rt.stop(sid)
            except Exception as ex:cleanup[sid]=dict(error=type(ex).__name__)
    states=rt.states()
    save('final.json',dict(error=error,states=states,cleanup=cleanup,image=rt.inspect()['Image'],
        models={n:rt.triton('v2/models/'+n+'/config') for n in MODELS},resource=rt.resources(),
        original_artifacts_preserved=all(file_hash(f)==h for f,h in p['model_files'].items()),
        presets_preserved=all(file_hash(f)==h for f,h in p['presets'].items()),env_preserved=file_hash(ENV)==p['env_hash']))
    print('Final failed=',bool(error),'active=',sum(bool(q.get('child_pid')) for q in states.values()),flush=True)
if __name__=='__main__':
    error=None
    try:prepare();reflect();contract();run()
    except Exception as e:error=dict(type=type(e).__name__,message=str(e),trace=traceback.format_exc());save('failure-'+str(time.time_ns())+'.json',error)
    finally:finish(error)