#!/usr/bin/env python3
"""Exclusive operational handoff: ASR on GPU1, or GPU1 released for future tests.
Never starts/stops ComfyUI, never kills unknown GPU processes, never prints API keys.
"""
import argparse,fcntl,json,pathlib,subprocess,time,urllib.request
ROOT=pathlib.Path(__file__).resolve().parent
CONTAINER='followup-audio-asr'

def command(args):
    p=subprocess.run(args,capture_output=True,text=True,timeout=700)
    if p.returncode: raise RuntimeError('COMMAND_FAILED: '+args[0])
    return p.stdout

def gpu_pids():
    text=command(['nvidia-smi','--id=1','--query-compute-apps=pid','--format=csv,noheader,nounits'])
    return {int(line.strip()) for line in text.splitlines() if line.strip()}

def container_state():
    p=subprocess.run(['docker','inspect','--format','{{json .State}}',CONTAINER],capture_output=True,text=True)
    if p.returncode:
        if 'no such object' in p.stderr.lower() or 'no such container' in p.stderr.lower():return {'Running':False,'Status':'absent'}
        raise RuntimeError('CONTAINER_STATE_UNKNOWN')
    return json.loads(p.stdout)

def owned_pids(state):
    if not state.get('Running'):return set()
    lines=command(['docker','top',CONTAINER,'-eo','pid']).splitlines()[1:]
    return {int(line.strip()) for line in lines if line.strip()}

def ensure_exclusive(active,owned):
    if active-owned: raise RuntimeError('GPU1_HAS_OTHER_WORKLOAD: refusing to stop or replace it')

def idle_values(body):
    values=[]
    for name in ['vllm:num_requests_running','vllm:num_requests_waiting']:
        rows=[float(l.rsplit(' ',1)[1]) for l in body.splitlines() if l.startswith(name+'{')]
        if not rows:raise RuntimeError('ASR_DRAIN_METRICS_MISSING')
        values.append(sum(rows))
    return values

def drain():
    zeros=0;deadline=time.monotonic()+600
    while time.monotonic()<deadline:
        with urllib.request.urlopen('http://127.0.0.1:8001/metrics',timeout=5) as r: values=idle_values(r.read().decode())
        zeros=zeros+1 if values==[0,0] else 0
        if zeros>=3:return
        time.sleep(5)
    raise RuntimeError('ASR_BUSY: unchanged; wait for requests to finish')

def main():
    p=argparse.ArgumentParser();p.add_argument('mode',choices=['status','asr','test']);args=p.parse_args()
    lock=open(ROOT/'.gpu1-mode.lock','a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    state=container_state();active=gpu_pids();owned=owned_pids(state)
    if args.mode=='status':
        print(json.dumps({'asr_status':state['Status'],'asr_running':state['Running'],'gpu1_processes':sorted(active),'other_gpu1_processes':sorted(active-owned),'test_service_started':False}));return
    ensure_exclusive(active,owned)
    compose=['docker','compose','--project-directory',str(ROOT),'--env-file',str(ROOT/'.env'),'-f',str(ROOT/'compose.yaml')]
    if args.mode=='test':
        if state['Running']:
            drain();ensure_exclusive(gpu_pids(),owned_pids(container_state()))
            command(compose+['stop','--timeout','660','asr'])
        if gpu_pids():raise RuntimeError('GPU1_NOT_FREE')
        result={'mode':'test','asr_stopped':True,'gpu1_free':True,'comfyui_test_started':False}
    else:
        if not state['Running']:
            if active:raise RuntimeError('GPU1_NOT_FREE')
            command(compose+['up','-d','--no-build','--pull','never','asr'])
        deadline=time.monotonic()+480
        while time.monotonic()<deadline:
            state=container_state()
            if state.get('Health',{}).get('Status')=='healthy':break
            if state['Status'] in ['exited','dead','restarting']:raise RuntimeError('ASR_START_FAILED')
            time.sleep(5)
        else:raise RuntimeError('ASR_READINESS_TIMEOUT')
        result={'mode':'asr','asr_healthy':True,'gpu1_processes':sorted(gpu_pids()),'comfyui_test_started':False}
    target=ROOT/'.gpu1-mode.json';stage=target.with_suffix('.tmp');stage.write_text(json.dumps(result)+'\n');stage.chmod(0o600);stage.replace(target)
    print(json.dumps(result))
if __name__=='__main__':main()