from ops import *
import sys,yaml
mode=sys.argv[1]
assert mode in ('on','on-r2','off','final')
overlay=PRIVATE/('freshness-'+mode+'.compose.yaml')
overlay.write_text(yaml.safe_dump({'services':{'analytics':{'environment':{
    'CONTEXTRA_PROVIDER_FRESHNESS_ENABLED':'false' if mode=='off' else 'true',
    'CONTEXTRA_PROVIDER_STARTUP_SECONDS':'10',
    'CONTEXTRA_METADATA_VISUAL_ENABLED':'false'}}}}))
overlay.chmod(0o600)
image=json.loads((HERE/'image-01df389ca6.json').read_text())['id']
assert not inspect('contextra-ai-gateway-1')['State']['Running']
rows=api('/api/v1/streams')['streams']
assert all((r['state']=='stopped' or mode=='final') and not r.get('cleanup_pending')
           and not r.get('child_pid') and all(not p.get('child_pid') for p in r.get('provider_processes',{}).values()) for r in rows)
deadline=time.monotonic()+60
while True:
    a=api('/api/v1/metadata/status')['artifacts']
    if not a.get('active') and not a.get('rendering'):break
    if time.monotonic()>deadline:raise RuntimeError('ARTIFACT_DRAIN')
    time.sleep(1)
full=compose()+['-f',str(overlay)]
if mode=='final':
    target=Path('/workspace/contextra-runtime-config/metadata-active/compose.yaml')
    shutil.copyfile(target,PRIVATE/'metadata-compose.before')
    config=yaml.safe_load(target.read_text())
    config['services']['analytics'].setdefault('environment',{}).update(
        yaml.safe_load(overlay.read_text())['services']['analytics']['environment'])
    target.write_text(yaml.safe_dump(config,sort_keys=False));target.chmod(0o600)
    full=compose()
expected=json.loads((PRIVATE/'compose-before.json').read_text())
expected['services']['analytics']['image']=image
expected['services']['analytics']['environment'].update(yaml.safe_load(overlay.read_text())['services']['analytics']['environment'])
assert json.loads(cmd(full+['config','--format','json']))==expected
save('reflection-'+mode+'-intent.json',dict(image=image,mode=mode,stopped=len(rows)))
assert shutil.disk_usage('/').free>15*1024**3
cmd(['docker','stop','--signal=SIGTERM','--timeout=-1','contextra-analytics-1'],timeout=90)
cmd(full+['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
save('reflection-'+mode+'-done.json',dict(image=image,mode=mode,ready=True))
print(json.dumps(dict(image=image,mode=mode,ready=True)),flush=True)