#!/usr/bin/env python3
"""Stop only existing router/controller/model; always restore the same containers."""
import json, pathlib, subprocess, sys, time, signal
ROOT=pathlib.Path('/data/benchmarks/medium-microbatch-20260920')
NAMES=['mike-ai-router','mike-ai-profile-controller','mike-ai-llama-medium']
def run(*args,check=True,timeout=90):
    return subprocess.run(args,capture_output=True,text=True,check=check,timeout=timeout)
def stop_signal(*_): raise RuntimeError('Supervisor interrupted')
signal.signal(signal.SIGTERM,stop_signal); signal.signal(signal.SIGINT,stop_signal)
# Refuse if the known production state has changed, or if requests are active.
for name in NAMES:
    assert run('docker','inspect',name,'--format','{{.State.Running}}').stdout.strip()=='true',name
for attempt in range(60):
    slots=json.loads(run('docker','exec',NAMES[-1],'curl','-fsS','http://127.0.0.1:8080/slots').stdout)
    if not any(s['is_processing'] for s in slots): break
    if attempt==0: print('WAIT production request active; no interruption',flush=True)
    time.sleep(3)
else: raise RuntimeError('Production remained busy for 180s; no services stopped')
assert not run('docker','ps','-q','--filter','name=^mike-ai-microbatch-test$').stdout.strip(),'Existing experiment'
production=json.loads(run('docker','inspect',NAMES[-1]).stdout)[0]
(ROOT/'production.json').write_text(json.dumps({'image':production['Image'],'args':production['Args']},indent=2)+'\n')
child=None
try:
    run('docker','stop','-t','30',*NAMES[:2])
    # Drain requests already handed to the model, before unloading it.
    for _ in range(120):
        slots=json.loads(run('docker','exec',NAMES[-1],'curl','-fsS','http://127.0.0.1:8080/slots').stdout)
        if not any(s['is_processing'] for s in slots): break
        time.sleep(2)
    else: raise RuntimeError('Model did not drain')
    run('docker','stop','-t','30',NAMES[-1])
    child=subprocess.Popen(['python3',str(ROOT/'run.py'),sys.argv[1]])
    code=child.wait(timeout=600)
    if code: raise RuntimeError('Benchmark failed: '+str(code))
finally:
    if child is not None and child.poll() is None:
        child.terminate()
        try: child.wait(timeout=20)
        except subprocess.TimeoutExpired: child.kill(); child.wait(timeout=10)
    run('docker','rm','-f','mike-ai-microbatch-test',check=False)
    run('docker','start',NAMES[-1])
    for _ in range(150):
        if run('docker','inspect',NAMES[-1],'--format','{{.State.Health.Status}}').stdout.strip()=='healthy': break
        time.sleep(2)
    else: raise RuntimeError('Restored Medium did not become healthy')
    run('docker','start',NAMES[1],NAMES[0])
    for _ in range(60):
        statuses=[run('docker','inspect',name,'--format','{{.State.Health.Status}}').stdout.strip() for name in NAMES[:2]]
        if all(status=='healthy' for status in statuses): break
        time.sleep(2)
    else: raise RuntimeError('Restored router/controller did not become healthy')
    print('RESTORED existing medium/controller/router; all healthy',flush=True)