from __future__ import annotations
import argparse, json, os, time
from pathlib import Path
import tensorrt as trt

LOGGER=trt.Logger(trt.Logger.INFO)

def build(onnx_path:Path, engine_path:Path, workspace_gib:int=8, fp16=True):
    parse_path = onnx_path
    # TensorRT 11 is strongly typed and removed BuilderFlag.FP16. NVIDIA's
    # supported migration path is ModelOpt AutoCast before engine building.
    if fp16 and not hasattr(trt.BuilderFlag, "FP16"):
        import onnx
        from modelopt.onnx.autocast import convert_to_mixed_precision
        mixed_path = onnx_path.with_name(onnx_path.stem + "_mixed_fp16.onnx")
        print(f"AUTOCAST {onnx_path.name} -> {mixed_path.name}", flush=True)
        converted = convert_to_mixed_precision(
            onnx_path=str(onnx_path), low_precision_type="fp16",
            keep_io_types=True, providers=["cuda:0", "cpu"],
            init_conversion_max_bytes=512 * 1024 * 1024,
            use_standalone_type_inference=True,
        )
        onnx.save(converted, str(mixed_path))
        parse_path = mixed_path
    builder=trt.Builder(LOGGER)
    flags = 0
    if hasattr(trt.NetworkDefinitionCreationFlag, 'EXPLICIT_BATCH'):
        flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network=builder.create_network(flags)
    parser=trt.OnnxParser(network, LOGGER)
    data=parse_path.read_bytes()
    if not parser.parse(data):
        errors='\n'.join(str(parser.get_error(i)) for i in range(parser.num_errors))
        raise RuntimeError(errors)
    config=builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_gib*(1<<30))
    if hasattr(config, 'builder_optimization_level'): config.builder_optimization_level=5
    if fp16 and hasattr(trt.BuilderFlag, "FP16"): config.set_flag(trt.BuilderFlag.FP16)
    # Engines are intentionally built for the current L4 and not marked HW-compatible.
    t=time.time(); serialized=builder.build_serialized_network(network, config)
    if serialized is None: raise RuntimeError('TensorRT build returned None')
    engine_path.parent.mkdir(parents=True,exist_ok=True); engine_path.write_bytes(serialized)
    return {'seconds':round(time.time()-t,3),'bytes':engine_path.stat().st_size,'fp16':fp16}

def main():
    ap=argparse.ArgumentParser(); ap.add_argument('--root',default='/tmp/cf-onnx-out'); ap.add_argument('--component',choices=['dinov2','dsine','all'],default='all'); a=ap.parse_args()
    root=Path(a.root); out={}; names=['dinov2','dsine'] if a.component=='all' else [a.component]
    for name in names:
        src=root/f'onnx/{name}/model.onnx'; dst=root/f'engines/l4-sm89/{name}.plan'
        print('BUILD',name,flush=True); out[name]=build(src,dst)
    import torch
    out['environment']={'gpu':torch.cuda.get_device_name(0),'cc':list(torch.cuda.get_device_capability(0)),'cuda':torch.version.cuda,'tensorrt':trt.__version__}
    meta=root/f'trt_build_{a.component}_meta.json'; meta.write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()