from ops import *
import sys,yaml
mode=sys.argv[1];assert mode in ('diagnostic','final')
assert not inspect('contextra-ai-gateway-1')['State']['Running']
deadline=time.monotonic()+90
while True:
    rows=api('/api/v1/streams')['streams'];a=api('/api/v1/metadata/status')['artifacts']
    assert all(r['state'] in ('stopped','error') and not r.get('child_pid') and not r.get('cleanup_pending') for r in rows)
    assert all(r['state']!='error' or (r.get('last_cleanup',{}).get('child_reaped') and r.get('last_cleanup',{}).get('writer_released')) for r in rows)
    if not a.get('active') and not a.get('rendering'):break
    if time.monotonic()>deadline:raise RuntimeError('ARTIFACT_DRAIN_FAILED')
    time.sleep(1)
image=json.loads((HERE/'image.json').read_text())['id']
args=compose()
if mode=='diagnostic':
    overlay={'services':{'analytics':{'image':image,'environment':{
        'STREAMS_CONFIG_PATH':'/compare/base.yaml','STREAMS_RUNTIME_PATH':'/compare/runtime.yaml',
        'CONTEXTRA_PIPELINE_DIAGNOSTICS':'true'},'volumes':[str(PRIVATE)+':/compare:rw']}}}
    path=PRIVATE/'compose-diagnostic.yaml';path.write_text(yaml.safe_dump(overlay));args+=['-f',str(path)]
expected=json.loads((PRIVATE/'compose-before.json').read_text())
actual=json.loads(cmd(args+['config','--format','json']))
check=json.loads(json.dumps(actual));check['services']['analytics']['image']=expected['services']['analytics']['image']
if mode=='diagnostic':
    for k in ('STREAMS_CONFIG_PATH','STREAMS_RUNTIME_PATH','CONTEXTRA_PIPELINE_DIAGNOSTICS'):
        if k in expected['services']['analytics']['environment']:check['services']['analytics']['environment'][k]=expected['services']['analytics']['environment'][k]
        else:check['services']['analytics']['environment'].pop(k,None)
    check['services']['analytics']['volumes']=[v for v in check['services']['analytics']['volumes'] if v['target']!='/compare']
assert check==expected,'UNEXPECTED_COMPOSE_DELTA'
save('transition-'+mode+'-intent.json',dict(image=image,stopped=len(rows)))
cmd(['docker','stop','--signal=SIGTERM','--timeout=-1','contextra-analytics-1'],timeout=90)
cmd(args+['up','-d','--no-deps','--no-build','--pull','never','analytics'],timeout=150)
deadline=time.monotonic()+120
while True:
    try:api('/ready');break
    except Exception:
        if time.monotonic()>deadline:raise RuntimeError('READY_FAILED')
        time.sleep(2)
assert inspect('contextra-analytics-1')['Image']==image
rows=api('/api/v1/streams')['streams'];assert all(r['state']=='stopped' for r in rows)
assert len(rows)==(35 if mode=='final' else 11)
assert hashlib.sha256(Path('/data/contextra/runtime-state/streams_runtime.yaml').read_bytes()).hexdigest()==json.loads((HERE/'config-fingerprints.json').read_text())['original']
save('transition-'+mode+'-done.json',dict(image=image,ready=True,channels=len(rows)))
print(json.dumps(dict(mode=mode,image=image,ready=True,channels=len(rows))),flush=True)