"""Build verified ONNX sequentially on AI2 GPU; successful candidates are not repeated."""
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys

STAGE=Path('/storage/contextra/deploy/ai2-20261006-r01')
SOURCE=Path('/workspace/contextra/source/4f0bd4e-ai2-tools-r01')
MODELS=Path('/storage/contextra/models')
sys.path.insert(0,str(SOURCE))
from deployment_kit import models
from deployment_kit.core import write_json,sha256

def cmd(args,timeout=30):
    return subprocess.check_output(args,text=True,stderr=subprocess.STDOUT,timeout=timeout)

def main():
    os.umask(0o077)
    assert cmd(['hostname']).strip()=='icertirnd-ai'
    assert not cmd(['nvidia-smi','--query-compute-apps=pid','--format=csv,noheader,nounits']).strip()
    assert shutil.disk_usage('/').free>=15*1024**3
    original_once=models._once
    def recorded_once(image,mounts,command,**kwargs):
        if kwargs.get('evidence') is None:
            output=next(Path(src) for src,dest,_ in mounts if dest=='/output')
            stage=command[2] if command[0]=='python3' else Path(command[0]).name
            kwargs['evidence']=output/(stage+'-container.log')
        return original_once(image,mounts,command,**kwargs)
    models._once=recorded_once
    receipt=STAGE/'tooling-ready.json'
    if receipt.exists():
        tool=json.loads(receipt.read_text())['image']
        assert json.loads(cmd(['docker','image','inspect',tool]))[0]['Id']==tool
    else:
        base=json.loads((STAGE/'images-resolved.json').read_text())['model-tooling']
        assert json.loads(cmd(['docker','image','inspect',base]))[0]['Id']==base
        context=STAGE/'model-tooling-build';context.mkdir(mode=0o700)
        shutil.copyfile(SOURCE/'deployment_kit/engine_contract.cpp',context/'engine_contract.cpp')
        (context/'Dockerfile').write_text('FROM '+base+'\nUSER root\nCOPY engine_contract.cpp /opt/contextra-model-tools/engine_contract.cpp\nRUN g++ -std=c++17 -O2 /opt/contextra-model-tools/engine_contract.cpp -I/usr/local/cuda/include -lnvinfer -lnvinfer_plugin -o /usr/local/bin/contextra-engine-contract\n')
        tag='contextra-model-tooling:ai2-4f0bd4e-20261006-probe-r01'
        with (STAGE/'model-tooling-build.log').open('x') as log:
            subprocess.run(['docker','build','--pull=false','--network=none','--memory=2g',
                '--cpu-period=100000','--cpu-quota=200000','-t',tag,str(context)],
                stdout=log,stderr=log,check=True,timeout=180,env=dict(os.environ,DOCKER_BUILDKIT='0'))
        tool=json.loads(cmd(['docker','image','inspect',tag]))[0]['Id']
        write_json(receipt,dict(image=tool,tag=tag,base=base,probe_source_sha256=sha256(context/'engine_contract.cpp')),new=True)
    result_path=STAGE/'models-result.json'
    results=json.loads(result_path.read_text()) if result_path.exists() else {}
    repo=MODELS/'triton';repo.mkdir(mode=0o755,parents=True,exist_ok=True)
    for name in ('medium','pose','fire','crowd'):
        if name in results:
            assert sha256(Path(results[name]['candidate'])/'model.plan')==results[name]['engine_sha256']
            print('ALREADY_VERIFIED '+name,flush=True)
            continue
        source=STAGE/'assets'/(name+'.onnx')
        assert sha256(source)==models.PROFILES[name]['source_onnx']
        suffix='r03' if name=='medium' else 'r01'
        candidate=MODELS/'candidates'/(name+'-batch8-ai2-20261006-'+suffix)
        assert not candidate.exists(), 'Review partial candidate before resuming '+name
        if name=='medium':
            prior=MODELS/'candidates/medium-batch8-ai2-20261006-r02'
            m=models.candidate_manifest(prior)
            assert m['status']=='CONTRACT_PASS' and m['smoke']=='NOT_RUN'
            candidate.mkdir(mode=0o700)
            for filename in ('manifest.json','profile.json','config.pbtxt'):
                shutil.copyfile(prior/filename,candidate/filename)
            os.link(prior/'model.plan',candidate/'model.plan')
            assert sha256(candidate/'model.plan')==m['engine_sha256']
            write_json(candidate/'engine-reuse.json',{'source':str(prior),'engine_sha256':m['engine_sha256'],'rebuild':False},new=True)
            print('MODEL_ENGINE_REUSED '+name,flush=True)
        else:
            print('MODEL_BUILD_START '+name,flush=True)
            m=models.build(name,source,candidate,tool,batch=8,instances=8,apply=True)
        print('MODEL_CONTRACT_PASS '+name,flush=True)
        smoke=models.verify_candidate(candidate,tool,apply=True)
        models.install_candidate(candidate,repo,apply=True,model_version=8)
        results[name]=dict(candidate=str(candidate),engine_sha256=m['engine_sha256'],smoke=smoke)
        write_json(result_path,results)
        print('MODEL_SMOKE_PASS '+name,flush=True)
    print('FOUR_BATCH8_MODELS_READY',flush=True)

if __name__=='__main__':main()