#!/usr/bin/env python3
"""
build_trt_engine.py — Convert ONNX model to TensorRT .engine file.

Run this ONCE on the Jetson Orin Nano Super to build the optimised engine.
The engine is hardware-specific and cannot be shared between GPU families.

Usage:
    python3 build_trt_engine.py --onnx yolov8n.onnx --engine yolov8n.engine
    python3 build_trt_engine.py --onnx yolov8n.onnx --engine yolov8n.engine --fp16
    python3 build_trt_engine.py --onnx yolov8n.onnx --engine yolov8n.engine --fp16 --batch 1

Alternatively, use the trtexec CLI tool (ships with JetPack):
    /usr/src/tensorrt/bin/trtexec \\
        --onnx=yolov8n.onnx \\
        --fp16 \\
        --saveEngine=yolov8n.engine \\
        --workspace=2048

Model download (YOLOv8n):
    pip3 install ultralytics
    python3 -c "from ultralytics import YOLO; YOLO('yolov8n.pt').export(format='onnx', imgsz=640)"

Model download (YOLOv5s):
    python3 -c "
    import torch
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
    model.export(format='onnx')
    "

Engine location:
    Place the built .engine file at the path specified by `engine_path`
    in person_detection_params.yaml (default: models/yolov8n.engine)
"""

import argparse
import sys
import os


def build_engine(onnx_path: str, engine_path: str, fp16: bool, batch_size: int,
                 workspace_mb: int) -> bool:
    """Build TensorRT engine from ONNX model."""
    try:
        import tensorrt as trt
    except ImportError:
        print('ERROR: tensorrt not found. Run on Jetson with JetPack installed.')
        print('       Alternatively use trtexec (see script header).')
        return False

    logger = trt.Logger(trt.Logger.VERBOSE)
    builder = trt.Builder(logger)
    network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(network_flags)
    parser = trt.OnnxParser(network, logger)

    print(f'[+] Parsing ONNX: {onnx_path}')
    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(f'    Parse error {i}: {parser.get_error(i)}')
            return False
    print(f'    Network inputs: {network.num_inputs}')
    print(f'    Network outputs: {network.num_outputs}')

    config = builder.create_builder_config()
    config.set_memory_pool_limit(
        trt.MemoryPoolType.WORKSPACE, workspace_mb * 1024 * 1024)

    if fp16 and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
        print('[+] FP16 mode enabled')
    elif fp16:
        print('[!] FP16 not supported on this platform, using FP32')

    # Dynamic batch profile (batch_size = 1 for real-time inference)
    profile = builder.create_optimization_profile()
    input_name = network.get_input(0).name
    input_shape = network.get_input(0).shape  # e.g., [-1, 3, 640, 640]
    min_shape = (1, input_shape[1], input_shape[2], input_shape[3])
    opt_shape = (batch_size, input_shape[1], input_shape[2], input_shape[3])
    max_shape = (batch_size, input_shape[1], input_shape[2], input_shape[3])
    profile.set_shape(input_name, min_shape, opt_shape, max_shape)
    config.add_optimization_profile(profile)

    print(f'[+] Building engine (this may take 5–15 minutes on first run)...')
    serialized = builder.build_serialized_network(network, config)
    if serialized is None:
        print('ERROR: Engine build failed.')
        return False

    os.makedirs(os.path.dirname(os.path.abspath(engine_path)), exist_ok=True)
    with open(engine_path, 'wb') as f:
        f.write(serialized)

    size_mb = os.path.getsize(engine_path) / (1024 * 1024)
    print(f'[+] Engine saved: {engine_path} ({size_mb:.1f} MB)')
    return True


def verify_engine(engine_path: str) -> bool:
    """Quick sanity check: deserialise and print binding info."""
    try:
        import tensorrt as trt
        import numpy as np
        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f, trt.Runtime(logger) as rt:
            engine = rt.deserialize_cuda_engine(f.read())
        print(f'\n[+] Engine verified: {engine.num_io_tensors} tensors')
        for i in range(engine.num_io_tensors):
            name = engine.get_tensor_name(i)
            shape = engine.get_tensor_shape(name)
            dtype = engine.get_tensor_dtype(name)
            mode = 'IN ' if engine.get_tensor_mode(name).name == 'INPUT' else 'OUT'
            print(f'    {mode} [{i}] {name}: {list(shape)} {dtype}')
        return True
    except Exception as e:
        print(f'[!] Engine verification failed: {e}')
        return False


def main():
    parser = argparse.ArgumentParser(
        description='Build TensorRT engine from ONNX model')
    parser.add_argument('--onnx', required=True,
                        help='Path to input ONNX model')
    parser.add_argument('--engine', required=True,
                        help='Path for output .engine file')
    parser.add_argument('--fp16', action='store_true', default=True,
                        help='Enable FP16 precision (default: True)')
    parser.add_argument('--no-fp16', dest='fp16', action='store_false',
                        help='Disable FP16, use FP32')
    parser.add_argument('--batch', type=int, default=1,
                        help='Batch size (default: 1)')
    parser.add_argument('--workspace', type=int, default=2048,
                        help='Builder workspace in MB (default: 2048)')
    parser.add_argument('--verify', action='store_true',
                        help='Verify engine after build')
    args = parser.parse_args()

    if not os.path.isfile(args.onnx):
        print(f'ERROR: ONNX model not found: {args.onnx}')
        print()
        print('Download YOLOv8n ONNX:')
        print('    pip3 install ultralytics')
        print('    python3 -c "from ultralytics import YOLO; '
              'YOLO(\'yolov8n.pt\').export(format=\'onnx\', imgsz=640)"')
        sys.exit(1)

    ok = build_engine(args.onnx, args.engine, args.fp16, args.batch, args.workspace)
    if not ok:
        sys.exit(1)

    if args.verify:
        verify_engine(args.engine)

    print('\nNext step: update person_detection_params.yaml:')
    print(f'    engine_path: "{os.path.abspath(args.engine)}"')


if __name__ == '__main__':
    main()
