import json
import os
import pathlib
import subprocess
import time

os.umask(0o077)
root = pathlib.Path(__file__).parent
name = 'codex-mux-copy-r01'
def gpu_used():
    return int(subprocess.check_output(['nvidia-smi', '--id=0', '--query-gpu=memory.used',
                '--format=csv,noheader,nounits'], text=True, timeout=5).strip())
baseline = gpu_used()
command = ['docker', 'run', '--rm', '--pull', 'never', '--name', name,
           '--network', 'none', '--gpus', 'device=0', '--cpus', '2', '--memory', '2g',
           '--ulimit', 'core=0', '-e', 'NVIDIA_VISIBLE_DEVICES=0',
           '-e', 'NVIDIA_DRIVER_CAPABILITIES=compute,video,utility',
           '-e', 'LD_LIBRARY_PATH=/opt/nvidia/deepstream/deepstream/lib:/usr/local/cuda/lib64:/usr/local/nvidia/lib64',
           '--mount', 'type=bind,src=/workspace/contextra-analytics-runtime-lifecycle-refactor,dst=/opt/contextra,readonly',
           '--mount', f'type=bind,src={root},dst=/evidence', '--entrypoint', 'python3',
           'sha256:e2137309461015890de735ccdfde452aaa4809ef58c107858c672c72762cd239', '/evidence/probe.py']
start = time.monotonic()
peak = baseline
reason = None
with (root/'container.log').open('w') as log:
    process = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT)
    try:
        while process.poll() is None:
            used = gpu_used()
            peak = max(peak, used)
            if used > baseline + 2048 or time.monotonic() - start > 120:
                reason = 'resource_or_time_cap'
                break
            time.sleep(1)
    finally:
        if process.poll() is None:
            subprocess.run(['docker', 'kill', name], capture_output=True, timeout=10)
        process.wait(timeout=15)
summary = {'baseline_gpu_mib': baseline, 'peak_gpu_mib': peak,
           'after_gpu_mib': gpu_used(), 'elapsed_seconds': round(time.monotonic()-start, 2),
           'container_returncode': process.returncode, 'abort_reason': reason}
(root/'resources.json').write_text(json.dumps(summary, indent=2))
print(json.dumps(summary))
if (root/'result.json').exists():
    print((root/'result.json').read_text())