import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time

ROOT=Path(__file__).resolve().parents[1]
os.umask(0o077)
stage=sys.argv[1]
assert stage in ('export','build','probe')
intent=ROOT/(stage+'.intent.json')
assert not intent.exists(), 'Existing attempt: inspect journal and partial output before any resume'
image=json.loads((ROOT/'manifest.intake.json').read_text())['exporter']['image']
name='contextra-awiros-'+stage+'-r01'
base=['docker','run','--rm','--name',name,'--pull=never','--network=none','--cpus=4','--memory=12g',
      '--memory-swap=12g','--pids-limit=256','--cap-drop=ALL','--security-opt=no-new-privileges',
      '--mount','type=bind,source='+str(ROOT)+',target=/work',
      '--mount','type=bind,source='+str(ROOT/'source')+',target=/work/source,readonly',
      '--mount','type=bind,source=/data/contextra/model-artifacts/tool-envs/medium-193-20260909,target=/envroot,readonly',
      '-e','OMP_NUM_THREADS=4','-e','OPENBLAS_NUM_THREADS=1','-e','MKL_NUM_THREADS=4',
      '-e','TMPDIR=/work/tmp','-e','HF_HOME=/work/cache','-e','XDG_CACHE_HOME=/work/cache','-w','/work']
for folder in ('tmp','cache','engine'): (ROOT/folder).mkdir(exist_ok=True)
timeout=900
if stage=='export':
    command=base+['--entrypoint','/envroot/venv/bin/python',image,'-B','/work/tools/export.py']
elif stage=='build':
    onnx=json.loads((ROOT/'onnx-contract.json').read_text())['path']
    command=base+['--gpus=device=0','--entrypoint','/usr/bin/trtexec',image,
        '--onnx='+onnx,'--saveEngine=/work/engine/model.plan','--fp16',
        '--inputIOFormats=fp16:chw','--outputIOFormats=fp32:chw,fp32:chw',
        '--minShapes=input:1x3x512x512','--optShapes=input:4x3x512x512','--maxShapes=input:4x3x512x512',
        '--memPoolSize=workspace:1024','--skipInference','--profilingVerbosity=detailed']
else:
    command=base+['--gpus=device=0','--entrypoint','/work/tools/probe',image,'/work/engine/model.plan']
    timeout=90
lowest=shutil.disk_usage('/').free
assert lowest>15*1024**3
with intent.open('x') as f:json.dump(dict(stage=stage,command=command,timeout_seconds=timeout),f,indent=2)
with (ROOT/(stage+'.log')).open('x') as log:
    p=subprocess.Popen(command,stdout=log,stderr=log)
    start=time.monotonic();reason=None
    while p.poll() is None:
        lowest=min(lowest,shutil.disk_usage('/').free)
        if lowest<15*1024**3 or time.monotonic()-start>timeout:
            reason='ROOT_CAPACITY' if lowest<15*1024**3 else 'TIMEOUT'
            subprocess.run(['docker','stop','--time=10',name],capture_output=True,timeout=25)
            p.wait(timeout=20)
            break
        time.sleep(2)
result=dict(stage=stage,returncode=p.returncode,reason=reason,elapsed_seconds=time.monotonic()-start,lowest_root_bytes=lowest)
with (ROOT/(stage+'.result.json')).open('x') as f:json.dump(result,f,indent=2)
print(json.dumps(result),flush=True)
raise SystemExit(0 if p.returncode==0 and reason is None else 1)