import fcntl
import traceback
from ops import *

CONFIG=Path('/workspace/contextra-runtime-config/test-run/medium8-fp16-sharing.yaml')
COMPOSE=['docker','compose','--env-file',str(ENV),'-f',str(REPO/'docker-compose.yml'),'-f',str(ACTIVE)]

def atomic(path,data,expected):
    assert file_hash(path)==expected
    st=path.stat();fd,tmp=tempfile.mkstemp(prefix='.sharing-',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():
    plan=read('plan.json');image=read('image.json');engine=read('engine.json')
    assert not TARGET.exists() and not CONFIG.exists()
    for path,digest in plan['persistent_hashes'].items():assert file_hash(path)==digest
    rt=Runtime(plan['old_config'],R)
    code="""import json
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/RFDETRMedium/config.pbtxt').read_text(),m)
assert m.max_batch_size==4 and m.instance_group[0].count==6 and m.input[0].data_type==pb.TYPE_FP32
m.version_policy.specific.versions[:]=[4];m.instance_group[0].count=8;m.input[0].data_type=pb.TYPE_FP16
assert all(o.data_type==pb.TYPE_FP32 for o in m.output)
print(text_format.MessageToString(m))
"""
    config=cmd(['docker','exec',rt.container,'python3','-c',code])
    (TARGET/'4').mkdir(parents=True,mode=0o700)
    shutil.copy2(R/'medium-fp16.plan',TARGET/'4/model.plan');(TARGET/'config.pbtxt').write_text(config)
    cfg=plan['old_config'];cfg['expected_image']=image['image'];cfg['model'].update(version=4,instances=8)
    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']))]
    CONFIG.write_text(yaml.safe_dump(cfg,sort_keys=False));CONFIG.chmod(0o600)
    save('prepared.json',dict(config=str(CONFIG),engine=engine,image=image,model_hash=file_hash(TARGET/'config.pbtxt')))

def reflect():
    plan=read('plan.json');rt=Runtime(plan['old_config'],R)
    assert not (R/'reflect-start.json').exists()
    save('reflect-start.json',dict(at=time.time(),state='Stop originals then reflect candidate; failure leaves targets stopped'))
    for row in plan['rows']:
        if row['original']:save('old-stop-'+row['stream_id']+'.json',rt.stop(row['stream_id']))
    assert all(q['state']=='stopped' for q in rt.states().values())
    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/RFDETRMedium':v['source']=str(TARGET)
    atomic(ACTIVE,yaml.safe_dump(active,sort_keys=False).encode(),plan['persistent_hashes'][str(ACTIVE)])
    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:
        current=rt.inspect()
        if current['State'].get('Health',{}).get('Status')=='healthy':break
        time.sleep(1)
    else:raise RuntimeError('Readiness deadline')
    assert current['Image']==read('image.json')['image']
    assert dict(x.split('=',1) for x in current['Config']['Env'])==dict(x.split('=',1) for x in plan['inspect']['Config']['Env'])
    for k in ('NanoCpus','CpuQuota','CpusetCpus','Memory','MemorySwap','PidMode','RestartPolicy'):
        assert current['HostConfig'][k]==plan['inspect']['HostConfig'][k]
    old={m['Destination']:m['Source'] for m in plan['inspect']['Mounts']};now={m['Destination']:m['Source'] for m in current['Mounts']}
    old['/opt/contextra/models/triton/RFDETRMedium']=str(TARGET);assert old==now
    files=[p for p in read('image.json')['hashes'] if p.startswith('app/')]
    for path in files:
        actual=cmd(['docker','exec',rt.container,'sha256sum','/opt/contextra/'+path]).split()[0]
        assert actual==read('image.json')['hashes'][path]
    save('reflected.json',dict(image=current['Image'],mounts=current['Mounts'],model=rt.triton('v2/models/RFDETRMedium/config')))

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('RFDETRMedium',[t],model_version='4',request_id='contract-'+str(b),client_timeout=30)
    assert r.get_response().id=='contract-'+str(b)
    entry={'batch':b,'outputs':{}}
    for name,shape in [('dets',(b,300,4)),('labels',(b,300,91))]:
        y=r.as_numpy(name);assert y.shape==shape and y.dtype==np.float32 and np.isfinite(y).all()
        entry['outputs'][name]={'shape':list(y.shape),'dtype':str(y.dtype),'finite':True}
    out.append(entry)
print(json.dumps(out))
"""
    value=json.loads(cmd(['docker','exec','contextra-analytics-1','python3','-c',code],150))
    save('contract.json',dict(batches=value,quality='not evaluated',request_identity=True))

class SharingRuntime(Runtime):
    def sample(self,rows):
        sample=super().sample(rows)
        sample['cuda_regions']=self.triton('v2/cudasharedmemory/status')
        assert len({r['name'] for r in sample['cuda_regions']})==len(sample['cuda_regions'])
        return sample

def run():
    cfg=load_config(CONFIG);evidence=R/'ascent';assert not evidence.exists()
    evidence.mkdir(mode=0o700);rt=SharingRuntime(cfg,evidence)
    with Path(cfg['lock_file']).open('a+') as lock:
        fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
        plan=preflight(rt);rt.save('plan.json',plan)
        rt.save('conditions.json',dict(config=cfg,source_head=read('plan.json')['head'],
            tools={p.name:file_hash(p) for p in (REPO/'tools').glob('stream_benchmark*.py')},candidate=read('prepared.json')))
        save('ascent-start.json',dict(at=time.time(),count=1,off_comparison=False))
        result=Benchmark(rt).run(plan)
    return result

def finish(failure):
    cfg=load_config(CONFIG) if CONFIG.exists() else load_config(OLD)
    rt=Runtime(cfg,R/'ascent' if (R/'ascent').exists() else R)
    rows=read('plan.json')['rows'];cleanup={}
    if failure:
        for row in reversed(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()
    region_sets=[]
    for _ in range(3):
        region_sets.append(rt.triton('v2/cudasharedmemory/status'));time.sleep(1)
    stuck=set.intersection(*(set(r['name'] for r in v) for v in region_sets))
    save('final.json',dict(failure=failure,states=states,cleanup=cleanup,cuda_region_samples=region_sets,
        stationary_region_names=sorted(stuck),image=rt.inspect()['Image'],model=rt.triton('v2/models/RFDETRMedium/config'),
        resource=rt.resources(),registry_preserved=file_hash(cfg['registration_file'])==read('plan.json')['registry_sha256'],
        env_preserved=file_hash(ENV)==read('plan.json')['persistent_hashes'][str(ENV)],
        sender_preserved=cmd(['python3',str(REPO/'tools/rtsp_test.py'),'list'])==read('plan.json')['sender'],
        redis_preserved=json.loads(cmd(['docker','inspect','contextra-redis-1']))[0]['Id']==read('plan.json')['redis_id']))
    print('Finished; failed=',bool(failure),'active=',sum(bool(q.get('child_pid')) for q in states.values()),flush=True)

if __name__=='__main__':
    error=None
    try:
        prepare();reflect();contract()
        if run():raise RuntimeError('Ascent failed; no automatic restore')
    except Exception as e:
        error=dict(type=type(e).__name__,message=str(e),trace=traceback.format_exc());save('failure.json',error)
    finally:
        finish(error)