import json
import os
import pathlib
import subprocess
import time
import urllib.request

os.umask(0o077)
R = pathlib.Path('/evidence')
result = {'model_loads': 1, 'probe_runs': 0}
def save():
    (R/'result.json').write_text(json.dumps(result, indent=2))
def stats():
    with urllib.request.urlopen('http://127.0.0.1:8000/v2/models/RFDETRMedium/versions/2/stats', timeout=3) as r:
        return json.load(r)
save()
start = time.monotonic()
with (R/'triton.log').open('w') as log:
    server = subprocess.Popen(['/opt/tritonserver/bin/tritonserver', '--model-repository=/models',
        '--model-control-mode=explicit', '--load-model=RFDETRMedium', '--exit-timeout=10',
        '--pinned-memory-pool-byte-size=268435456', '--cuda-memory-pool-byte-size=0:67108864'],
        stdout=log, stderr=subprocess.STDOUT)
    try:
        ready = False
        while time.monotonic()-start < 120 and server.poll() is None:
            try:
                with urllib.request.urlopen('http://127.0.0.1:8000/v2/models/RFDETRMedium/versions/2/ready',timeout=2) as r:
                    ready = r.status == 200
                if ready:
                    break
            except Exception:
                pass
            time.sleep(1)
        result.update(model_ready=ready, model_load_s=time.monotonic()-start)
        save()
        if ready:
            (R/'stats-before.json').write_text(json.dumps(stats()))
            result['probe_runs'] = 1
            save()
            env = dict(os.environ, OPENBLAS_NUM_THREADS='1')
            with (R/'probe.log').open('w') as out:
                try:
                    result['probe_exit'] = subprocess.run(['python3', '/evidence/probe.py'], env=env,
                        stdout=out, stderr=subprocess.STDOUT, timeout=115).returncode
                except subprocess.TimeoutExpired:
                    result['probe_timeout'] = True
            (R/'stats-after.json').write_text(json.dumps(stats()))
    finally:
        server.terminate()
        try:
            server.wait(15)
            result['triton_killed'] = False
        except subprocess.TimeoutExpired:
            server.kill()
            server.wait(5)
            result['triton_killed'] = True
        result['triton_exit'] = server.returncode
        save()