import json
import os
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path

from model_tooling.manifests import fingerprints
from model_tooling.medium import checked_result, environment
from model_tooling.numerical import prepare, compare

os.umask(0o077)
run = Path('/data/contextra/model-artifacts/builds/medium-193-20260909-fp32-control-r01')
identity = json.loads((run / 'export/result.json').read_text())['input_identity']
onnx = checked_result(run, 'export', identity)
assert fingerprints(onnx)['sha256'] == '6ab93d8f1c64c369b2273b91738d554ce76466fb31f41f5fabb1f61d24a91999'
directory = run / 'build'
directory.mkdir(mode=0o700)
artifact = directory / 'model.plan'
tool = Path('/usr/bin/trtexec')
assert fingerprints(tool)['sha256'] == '507fcc1670bcac695eaaefaf1639e81eef40f6dd195eedfdf23b25be66aa2021'
command = [str(tool), '--onnx=' + str(onnx), '--saveEngine=' + str(artifact),
           '--noTF32', '--minShapes=input:1x3x576x576', '--optShapes=input:1x3x576x576',
           '--maxShapes=input:1x3x576x576', '--memPoolSize=workspace:1024', '--skipInference']
result = {'stage': 'build', 'status': 'FAILED', 'input_identity': identity,
          'started_at': datetime.now(timezone.utc).isoformat(), 'environment': environment(),
          'command': command, 'trtexec_sha256': fingerprints(tool)['sha256'],
          'experiment_script_sha256': fingerprints(Path(__file__))['sha256'],
          'precision': 'FP32_ENABLED_NO_TF32_NO_FP16', 'runtime_validation': 'NOT_RUN',
          'timeout_seconds': 1800, 'workspace_mib': 1024}
start = time.monotonic()
try:
    with (directory / 'stage.log').open('xb') as log:
        subprocess.run(command, check=True, timeout=1800, stdout=log, stderr=log)
    result['artifact'] = {'filename': artifact.name, 'relative_path': 'build/model.plan',
                          'fingerprint': fingerprints(artifact)}
    result['status'] = 'PASS'
except Exception as error:
    result['error_type'] = type(error).__name__
    raise
finally:
    result['elapsed_seconds'] = time.monotonic() - start
    result['finished_at'] = datetime.now(timezone.utc).isoformat()
    with (directory / 'result.json').open('x') as output:
        json.dump(result, output, indent=2)
print('build PASS', round(result['elapsed_seconds'], 2), flush=True)
prepare(run)
with (run / 'numerical/probe.log').open('xb') as log:
    subprocess.run(['/probe/engine_probe', str(artifact), str(run / 'numerical/input.bin'),
                    str(run / 'numerical')], check=True, timeout=290, stdout=log, stderr=log)
print('probe completed: one execution, warmup zero', flush=True)
raise SystemExit(0 if compare(run) else 2)