import json
import pathlib
import subprocess
import sys

PREP = pathlib.Path('/data/codex-artifacts/contextra/action-model-eval/2026/2026-10-26_contextra-camera24-prep_r01')
sys.path.insert(0, str(PREP/'helpers'))
from remote import run, SSH

LOCAL = pathlib.Path(__file__).resolve().parent.parent
REMOTE = '/storage/contextra/experiments/action-model-eval/2026/'+LOCAL.name
PREMOTE = '/storage/contextra/experiments/action-model-eval/2026/'+PREP.name
IMAGE = 'contextra-action-eval:20261026-camera24-r03'

def checked(args, **kw):
    p = run(args, **kw)
    if p.returncode:
        raise RuntimeError(p.stderr.decode()[-3000:] + p.stdout.decode()[-3000:])
    return p.stdout

def snapshot(label):
    data = checked(['docker','exec','-i','contextra-ai2-analytics-1','python3','-'],
                   data=(PREP/'helpers/snapshot.py').read_bytes(),timeout=90)
    (LOCAL/'reports'/('production_'+label+'.json')).write_bytes(data)
    d = json.loads(data)
    print(json.dumps(dict(label=label,health=d['health'],config_sha256=d['config_sha256'],
                         state_counts={s:sum(v['state']==s for v in d['streams'].values())
                                       for s in set(v['state'] for v in d['streams'].values())})))

def upload():
    checked(['mkdir','-p',REMOTE+'/helpers',REMOTE+'/reports',REMOTE+'/outputs',REMOTE+'/inputs',REMOTE+'/pose_models/RFDETRKeypointPreview/8'],sudo=True)
    checked(['chown','si:si',REMOTE,REMOTE+'/helpers',REMOTE+'/reports',REMOTE+'/outputs',REMOTE+'/inputs',REMOTE+'/pose_models',REMOTE+'/pose_models/RFDETRKeypointPreview',REMOTE+'/pose_models/RFDETRKeypointPreview/8'],sudo=True)
    checked(['chmod','700',REMOTE])
    for p in sorted((LOCAL/'helpers').glob('*.py')):
        r = subprocess.run(['scp','-q','-i',SSH[2],'-o',SSH[4],'-o','StrictHostKeyChecking=yes',
                            '-P','1022',str(p),'si@58.233.207.31:'+REMOTE+'/helpers/'+p.name],timeout=30)
        r.check_returncode()

def trial(stage,timeout=1200):
    free = int(checked(['nvidia-smi','--query-gpu=memory.free','--format=csv,noheader,nounits']).decode().strip())
    if free < 8192:
        raise RuntimeError('Insufficient GPU headroom; production untouched')
    name = 'contextra-action-eval-'+stage+'-20261026-r01'
    args = ['docker','run','--rm','--name',name,'--gpus','device=0','--user','0:0',
            '--read-only','--cap-drop','ALL','--cap-add','DAC_OVERRIDE','--security-opt','no-new-privileges',
            '--cpus','2','--memory','8g','--pids-limit','256','--network','none','--shm-size','256m',
            '--tmpfs','/tmp:rw,size=512m','-e','MPLCONFIGDIR=/tmp/mpl','-e','PYTHONDONTWRITEBYTECODE=1',
            '-e','OMP_NUM_THREADS=2','-e','OPENBLAS_NUM_THREADS=2',
            '-v',REMOTE+':/eval','-v',PREMOTE+':/work:ro']
    if stage == 'render':
        args.remove('--gpus'); args.remove('device=0')
        image = checked(['docker','inspect','contextra-ai2-analytics-1','--format','{{.Image}}']).decode().strip()
        args += ['--entrypoint','python3',image,'/eval/helpers/render.py']
    elif stage == 'pose':
        image = checked(['docker','inspect','contextra-ai2-analytics-1','--format','{{.Image}}']).decode().strip()
        args += ['-v','/storage/contextra/models/triton/RFDETRKeypointPreview/8:/eval/pose_models/RFDETRKeypointPreview/8:ro',
                 '--entrypoint','python3',image,'/eval/helpers/pose_extract.py']
    else:
        args += ['--entrypoint','python',IMAGE,'/eval/helpers/evaluate.py',stage]
    print('START',stage,flush=True)
    try:
        p = run(args,timeout=timeout)
    except BaseException:
        run(['docker','stop','--time','10',name],timeout=30)
        raise
    logfile=LOCAL/'reports'/(stage+'_console.log')
    if logfile.exists():
        logfile=LOCAL/'reports'/(stage+'_console_attempt'+str(len(list((LOCAL/'reports').glob(stage+'_console*')))+1)+'.log')
    logfile.write_bytes(p.stdout+p.stderr)
    print(p.stdout.decode()[-5000:],p.stderr.decode()[-3000:],flush=True)
    print('EXIT',p.returncode,flush=True)
    if p.returncode:
        raise SystemExit(p.returncode)

if __name__ == '__main__':
    (LOCAL/'reports').mkdir(parents=True,exist_ok=True)
    if sys.argv[1]=='upload': upload()
    elif sys.argv[1]=='snapshot': snapshot(sys.argv[2])
    else: trial(sys.argv[1])