from ops import *
import sys
phase=sys.argv[1];assert phase in ('a','c','final')
assert not inspect('contextra-icerti-ai-gateway-1')['State']['Running']
rows=api('/api/v1/streams')['streams']
assert all(r['state'] in (('stopped','error') if phase=='final' else ('stopped',)) and not r.get('cleanup_pending') for r in rows)
until=time.monotonic()+140
while time.monotonic()<until:
    a=api('/api/v1/metadata/status')['artifacts']
    t=training().get('status',{}).get('counts',{})
    if not a['active'] and not a['ledger'].get('pending') and not a['ledger']['storage']['reserved_jobs'] and not any(t.get(k) for k in ('CAPTURING','QUEUED','RUNNING')):break
    time.sleep(2)
else:raise RuntimeError('EVENT_DRAIN_BOUNDARY')
save('drain-'+phase+'.json',dict(ledger=ledger(),artifacts=a,training=training()))
doc=json.loads((PRIVATE/'compose-before.json').read_text())
flags={} if phase=='final' else dict(CONTEXTRA_TEST_NO_CAPTURE=str(phase=='c').lower(),CONTEXTRA_INPUT_TIMING_SECONDS='2400',
    CONTEXTRA_PIPELINE_DIAGNOSTICS='true',CONTEXTRA_INPUT_TIMELINE_LIMIT='2')
doc['services']['analytics']['environment'].update(flags)
path=(Path(json.loads((BASE/'prepared.json').read_text())['transaction'])/'compose.json'
      if phase=='final' else PRIVATE/('compose-'+phase+'.json'))
if phase!='final':save(path.name,doc,True)
cmd(['docker','stop','--time','90',CONTAINER],100)
state=inspect(CONTAINER)['State'];assert state['ExitCode']==0 and not state['OOMKilled']
cmd(['docker','compose','-p','contextra-icerti','-f',str(path),'up','-d','--no-deps','--no-build','--pull','never','analytics'],150)
until=time.monotonic()+240
while time.monotonic()<until:
    state=inspect(CONTAINER)['State'];assert state['Running'] and not state['OOMKilled']
    try:api('/ready');break
    except Exception:time.sleep(2)
else:raise RuntimeError('REFLECTION_READINESS')
env=dict(v.split('=',1) for v in inspect(CONTAINER)['Config']['Env'])
assert all(env.get(k)==v for k,v in flags.items())
assert inspect(CONTAINER)['Image']==doc['services']['analytics']['image']
save('reflect-'+phase+'.json',dict(image=inspect(CONTAINER)['Image'],flags=flags,ready=True))
print(json.dumps(dict(phase=phase,reflected=True)),flush=True)