from ops import *
import sys,yaml
mode=sys.argv[1]
assert mode in ('old','current','restore')
if inspect('contextra-ai-gateway-1')['State']['Running']:
    cmd(['docker','stop','--signal=SIGTERM','--timeout=-1','contextra-ai-gateway-1'],timeout=90)
for r in api('/api/v1/streams')['streams']:
    if r['state']!='stopped' or r.get('child_pid') or r.get('cleanup_pending'):
        request('/api/v1/streams/'+r['stream_id']+'/stop',{},timeout=90)
deadline=time.monotonic()+90
while True:
    rows=api('/api/v1/streams')['streams'];a=api('/api/v1/metadata/status')['artifacts']
    if all(r['state']=='stopped' and not r.get('child_pid') and not r.get('cleanup_pending') for r in rows) and not a.get('active') and not a.get('rendering'):break
    if time.monotonic()>deadline:raise RuntimeError('STOP_OR_ARTIFACT_DRAIN_FAILED')
    time.sleep(1)
image=(json.loads((HERE/'baseline.json').read_text())['containers']['contextra-analytics-1']['image']
       if mode=='restore' else json.loads((HERE/('image-'+mode+'.json')).read_text())['id'])
args=compose()
if mode!='restore':
    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-'+mode+'.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']))
if mode=='restore':assert actual==expected
else:
    check=json.loads(json.dumps(actual));check['services']['analytics']['image']=expected['services']['analytics']['image']
    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)==(len(json.loads((HERE/'baseline.json').read_text())['states']) if mode=='restore' else 10)
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)