"""Run all 60 rollouts exclusively on the eai host; resume completed episodes."""
import datetime
import json
import os
from pathlib import Path
import signal
import socket
import subprocess
import time

ROOT=Path(__file__).resolve().parent
MODELS=ROOT.parent/'b1k_sft'
VARIANTS=['human','codegen_success','mix_h50_c50']
PORT=18080
EXPECTED_HOST='eai-Raider-16-HX-AI-A2XWIG'
POLICY_PY=ROOT/'.policy-venv/bin/python'
SIM_PY=Path('/home/eai/miniconda3/envs/behavior-392dev/bin/python')
B1K=Path('/home/eai/code/b1k')
status={'host':socket.gethostname(),'started':datetime.datetime.now().isoformat(),'state':'starting','variants':{}}
children=[]

def save():
    status['updated']=datetime.datetime.now().isoformat()
    p=ROOT/'status.tmp';p.write_text(json.dumps(status,indent=2));p.replace(ROOT/'status.json')

def stop(p):
    if p and p.poll() is None:
        os.killpg(p.pid,signal.SIGTERM)
        try:p.wait(timeout=30)
        except subprocess.TimeoutExpired:
            os.killpg(p.pid,signal.SIGKILL);p.wait()

def missing(variant):
    root=ROOT/'results'/variant/'json'
    ids=[]
    for i in range(20):
        p=root/f'turning_on_radio_{301+i}_0.json'
        if not p.exists():ids.append(i);continue
        r=json.loads(p.read_text())
        assert r['instance_id']==301+i and r['task']=='turning_on_radio'
        assert 'success' in r and 'final' in r['q_score']
    return ids

def main():
    assert socket.gethostname()==EXPECTED_HOST,'Evaluation is restricted to eai'
    policy_env=os.environ.copy()
    policy_env.update(PYTHONPATH=':'.join(map(str,[ROOT,ROOT/'behavior-1k-solution/src',ROOT/'behavior-1k-solution/openpi/src',ROOT/'behavior-1k-solution/openpi/packages/openpi-client/src'])),CUDA_VISIBLE_DEVICES='0',XLA_PYTHON_CLIENT_PREALLOCATE='false',XLA_PYTHON_CLIENT_ALLOCATOR='platform',XLA_FLAGS='--xla_gpu_enable_triton_gemm=false',JAX_COMPILATION_CACHE_DIR=str(ROOT/'jax-cache'),JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS='1',TOKENIZERS_PARALLELISM='false',OMP_NUM_THREADS='4')
    sim_env=os.environ.copy()
    sim_env.update(PYTHONPATH=':'.join(map(str,[B1K/'OmniGibson',B1K/'bddl3',B1K/'joylo'])),OMNIGIBSON_HEADLESS='1',OMNIGIBSON_DATA_PATH=str(B1K/'datasets'),CUDA_VISIBLE_DEVICES='0',OMP_NUM_THREADS='4',PYTHONUNBUFFERED='1')
    for key in ['ISAAC_PATH','EXP_PATH','CARB_APP_PATH']:
        sim_env.pop(key,None)
    for variant in VARIANTS:
        directory=ROOT/'results'/variant;directory.mkdir(parents=True,exist_ok=True)
        remaining=missing(variant)
        if not remaining:
            status['variants'][variant]={'state':'complete','completed':20};save();continue
        ready=directory/'policy-ready.json';ready.unlink(missing_ok=True)
        status.update(state='loading_policy',active_variant=variant);save()
        with (directory/'policy.log').open('a') as log:
            policy=subprocess.Popen([str(POLICY_PY),str(ROOT/'serve_checkpoint.py'),'--variant',str(MODELS/variant),'--port',str(PORT),'--log',str(directory/'policy-events.jsonl'),'--ready',str(ready)],env=policy_env,stdout=log,stderr=subprocess.STDOUT,start_new_session=True)
            children.append(policy)
            try:
                deadline=time.monotonic()+1200
                while not ready.exists():
                    if policy.poll() is not None:raise RuntimeError(f'{variant} policy exited: {policy.returncode}')
                    if time.monotonic()>deadline:raise TimeoutError('Policy startup timeout')
                    time.sleep(2)
                for attempt in range(1,4):
                    remaining=missing(variant)
                    if not remaining:break
                    status.update(state='evaluating',active_variant=variant,attempt=attempt)
                    status['variants'][variant]={'state':'running','completed':20-len(remaining),'pending_indices':remaining};save()
                    cmd=[str(SIM_PY),str(ROOT/'run_simulator.py'),'--task-name','turning_on_radio','--mode','public_test','--host','127.0.0.1','--port',str(PORT),'--instance-indices',*map(str,remaining),'--output-dir',str(directory)]
                    with (directory/f'simulator-attempt{attempt}.log').open('a') as simlog:
                        sim=subprocess.Popen(cmd,env=sim_env,stdout=simlog,stderr=subprocess.STDOUT,start_new_session=True);children.append(sim)
                        while sim.poll() is None:
                            todo=missing(variant)
                            status['variants'][variant]={'state':'running','completed':20-len(todo),'pending_indices':todo};save()
                            if policy.poll() is not None:
                                stop(sim);raise RuntimeError(f'Policy failed during {variant} evaluation')
                            time.sleep(10)
                    status['last_simulator_exit']=sim.returncode;save()
                if missing(variant):raise RuntimeError(f'{variant} still missing indices {missing(variant)} after retries')
                status['variants'][variant]={'state':'complete','completed':20};save()
            finally:stop(policy)
    status['state']='complete';save()
    subprocess.run([str(SIM_PY),str(ROOT/'summarize.py')],check=True)

if __name__=='__main__':
    try:main()
    except BaseException as exc:
        status.update(state='failed',error=repr(exc));save();raise
    finally:
        for p in children:stop(p)