import json
import os
import pathlib
import subprocess
import sys
from preflight import OLD_IMAGE

os.umask(0o077)
R = pathlib.Path(__file__).parent
mode = sys.argv[1]
assert mode in ('source','image')
image = OLD_IMAGE if mode=='source' else json.loads((R/'build.json').read_text())['image_id']
output = R/mode
output.mkdir(mode=0o700)
source = pathlib.Path('/workspace/contextra-analytics-runtime-lifecycle-refactor') if mode=='source' else R/'snapshot'
name = 'codex-receiver-'+mode
args = ['docker','run','--rm','--pull','never','--runtime','runc','--network','none','--name',name,
        '--cap-drop','MKNOD','--cpus','2','--memory','4g','--ulimit','core=0',
        '-e','NVIDIA_VISIBLE_DEVICES=void','-e','PYTHONDONTWRITEBYTECODE=1',
        '-e','DATABASE_URL=postgresql://u:p@localhost:1/db','-e','REDIS_STREAMS_ENABLED=false',
        '-e','PYTHONPATH='+('/source:/opt/contextra' if mode=='source' else '/opt/contextra'),
        '-e','LD_LIBRARY_PATH=/driver:/opt/tritonserver/lib:/usr/local/cuda/lib64:/usr/lib/x86_64-linux-gnu']
for src,dst,ro in [(output,'/evidence',False),(source,'/source',True),(R/'verify.py','/verify.py',True),
                  (pathlib.Path('/usr/lib/x86_64-linux-gnu/libcuda.so.595.84'),'/driver/libcuda.so.1',True)]:
    args += ['--mount',f'type=bind,src={src},dst={dst}'+(',readonly' if ro else '')]
args += ['--entrypoint','timeout',image,'90s','python3','/verify.py']
try:
    p = subprocess.run(args,capture_output=True,text=True,timeout=100)
    (output/'console.log').write_text(p.stdout+p.stderr)
    print(p.stdout[-2000:])
    print('exit',p.returncode)
    assert p.returncode == 0
finally:
    subprocess.run(['docker','rm','-f',name],capture_output=True,timeout=10)