#!/usr/bin/env python3
"""
build_midas_trt_engine.py -- Build TensorRT FP16 engine for MiDaS_small from ONNX.

Usage:
    python3 build_midas_trt_engine.py \
        --onnx /mnt/nvme/saltybot/models/midas_small.onnx \
        --engine /mnt/nvme/saltybot/models/midas_small.engine \
        --fp16

Requires: tensorrt, pycuda
"""

import argparse
import os
import sys


def build_engine(onnx_path: str, engine_path: str, fp16: bool = True):
    try:
        import tensorrt as trt
    except ImportError:
        print('ERROR: tensorrt not found. Install TensorRT first.')
        sys.exit(1)

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

    print(f'Parsing ONNX model: {onnx_path}')
    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(f'  ONNX parse error: {parser.get_error(i)}')
            sys.exit(1)

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

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

    print('Building TensorRT engine (this may take several minutes)...')
    engine_bytes = builder.build_serialized_network(network, config)
    if engine_bytes is None:
        print('ERROR: Failed to build engine')
        sys.exit(1)

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

    size_mb = len(engine_bytes) / (1024 * 1024)
    print(f'Engine saved: {engine_path} ({size_mb:.1f} MB)')


def main():
    parser = argparse.ArgumentParser(
        description='Build TensorRT FP16 engine for MiDaS_small')
    parser.add_argument('--onnx', required=True,
                        help='Path to MiDaS ONNX model')
    parser.add_argument('--engine', required=True,
                        help='Output TRT engine path')
    parser.add_argument('--fp16', action='store_true', default=True,
                        help='Enable FP16 (default: True)')
    parser.add_argument('--fp32', action='store_true',
                        help='Force FP32 (disable FP16)')
    args = parser.parse_args()

    fp16 = not args.fp32
    build_engine(args.onnx, args.engine, fp16=fp16)


if __name__ == '__main__':
    main()
