"""Runs INSIDE an approved model-tooling image, never installs Python packages."""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import sys
import urllib.request


def onnx_contract(path, profile):
    import onnx
    model=onnx.load(path, load_external_data=False)
    if any(x.data_location==onnx.TensorProto.EXTERNAL for x in model.graph.initializer):
        raise ValueError('ONNX_EXTERNAL_DATA_REQUIRES_SEPARATE_PACKAGE')
    onnx.checker.check_model(model)
    model=onnx.shape_inference.infer_shapes(model)
    def desc(v):
        t=v.type.tensor_type
        dims=[d.dim_value if d.HasField('dim_value') else -1 for d in t.shape.dim]
        return {'dtype':t.elem_type,'shape':dims}
    tensors={x.name:desc(x) for x in (*model.graph.input,*model.graph.output)}
    r=profile['resolution']
    expected={'input':[-1,3,r,r], **{k:[-1,*v] for k,v in profile['outputs'].items()}}
    if set(tensors)!=set(expected): raise ValueError('ONNX_NAMES')
    for name,shape in expected.items():
        actual=tensors[name]['shape']
        # Some Pose output axes remain symbolic after ONNX inference. Exact output sizes
        # are mandatory again after TensorRT deserialization and batch execution.
        compatible=(actual==shape if name=='input' else len(actual)==len(shape)
                    and actual[0]==-1 and all(a==-1 or a==b for a,b in zip(actual[1:],shape[1:])))
        if not compatible or tensors[name]['dtype']!=onnx.TensorProto.FLOAT:
            raise ValueError('ONNX_SHAPE_DTYPE')
    unresolved=[name for name,value in tensors.items() if name!='input' and -1 in value['shape'][1:]]
    return {'onnx_version':onnx.__version__,'tensors':tensors,'exact_output_axes_deferred_to_engine':unresolved}


def engine_contract(path, profile):
    if os.environ.get('CONTEXTRA_ENGINE_PROBE'):
        import subprocess
        result=subprocess.run([os.environ['CONTEXTRA_ENGINE_PROBE'],str(path),profile['key'],str(profile['batch'])],
                              capture_output=True,text=True,timeout=50,check=True)
        return json.loads(result.stdout)
    import tensorrt as trt
    if not trt.__version__.startswith('10.16.'):
        raise ValueError('TENSORRT_VERSION_REQUIRES_REVIEW')
    logger=trt.Logger(trt.Logger.ERROR)
    with trt.Runtime(logger) as runtime:
        engine=runtime.deserialize_cuda_engine(Path(path).read_bytes())
        if engine is None: raise ValueError('ENGINE_DESERIALIZE')
        r=profile['resolution']
        expected={'input':[-1,3,r,r], **{k:[-1,*v] for k,v in profile['outputs'].items()}}
        actual={}
        for i in range(engine.num_io_tensors):
            name=engine.get_tensor_name(i)
            actual[name]={'shape':list(engine.get_tensor_shape(name)), 'dtype':str(engine.get_tensor_dtype(name)),
                          'mode':str(engine.get_tensor_mode(name))}
        if set(actual)!=set(expected): raise ValueError('ENGINE_NAMES')
        for name,shape in expected.items():
            dtype=trt.float16 if name=='input' else trt.float32
            mode=trt.TensorIOMode.INPUT if name=='input' else trt.TensorIOMode.OUTPUT
            if list(engine.get_tensor_shape(name))!=shape or engine.get_tensor_dtype(name)!=dtype or engine.get_tensor_mode(name)!=mode:
                raise ValueError('ENGINE_CONTRACT')
        bounds=[list(s) for s in engine.get_tensor_profile_shape('input',0)]
        if bounds!=[[b,3,r,r] for b in (1,profile['batch'],profile['batch'])]:
            raise ValueError('ENGINE_BATCH_PROFILE')
        return {'tensorrt_version':trt.__version__,'tensors':actual,'batch_bounds':bounds}


def smoke(url,profile):
    import numpy as np
    import tritonclient.grpc as client
    # The pinned Runtime ships gRPC; no gevent/http extra installation is needed.
    if url != 'http://127.0.0.1:8000': raise ValueError('ISOLATED_LOOPBACK_ONLY')
    c=client.InferenceServerClient(url='127.0.0.1:8001')
    results=[]
    try:
        for batch in range(1,profile['batch']+1):
            data=np.zeros((batch,3,profile['resolution'],profile['resolution']),dtype=np.float16)
            inp=client.InferInput('input',data.shape,'FP16'); inp.set_data_from_numpy(data)
            rid='deployment-contract-'+str(batch)
            result=c.infer(profile['name'],[inp],request_id=rid,client_timeout=5)
            if result.get_response().id!=rid: raise ValueError('REQUEST_ID')
            for name,tail in profile['outputs'].items():
                arr=result.as_numpy(name)
                if arr is None or list(arr.shape)!=[batch,*tail] or arr.dtype!=np.float32 or not np.isfinite(arr).all():
                    raise ValueError('INFERENCE_OUTPUT_CONTRACT')
            results.append(batch)
    finally: c.close()
    return {'batches':results,'shape_dtype_finite':'PASS','accuracy':'NOT_RUN','performance':'NOT_RUN'}


def main():
    p=argparse.ArgumentParser(); p.add_argument('action',choices=['onnx','engine','smoke','ready'])
    p.add_argument('--profile'); p.add_argument('--path'); p.add_argument('--url',default='http://127.0.0.1:8000')
    a=p.parse_args()
    profile=json.loads(Path(a.profile).read_text()) if a.profile else None
    if a.action=='ready':
        with urllib.request.urlopen(a.url+'/v2/models/'+profile['name']+'/ready',timeout=2) as r:
            if r.status!=200: raise ValueError('NOT_READY')
        result={'ready':True}
    elif a.action=='onnx': result=onnx_contract(a.path,profile)
    elif a.action=='engine': result=engine_contract(a.path,profile)
    else: result=smoke(a.url,profile)
    print(json.dumps(result,sort_keys=True))

if __name__=='__main__':
    try: main()
    except Exception as error:
        known={'ONNX_EXTERNAL_DATA_REQUIRES_SEPARATE_PACKAGE','ONNX_NAMES','ONNX_SHAPE_DTYPE',
               'TENSORRT_VERSION_REQUIRES_REVIEW','ENGINE_DESERIALIZE','ENGINE_NAMES','ENGINE_CONTRACT',
               'ENGINE_BATCH_PROFILE','REQUEST_ID','INFERENCE_OUTPUT_CONTRACT','ISOLATED_LOOPBACK_ONLY','NOT_READY'}
        print(json.dumps({'error':type(error).__name__,'stage':'MODEL_PROBE_FAILED',
                          'code':str(error) if str(error) in known else 'PROBE_EXECUTION_FAILED'}))
        sys.exit(2)