#!/usr/bin/env python3
"""
build_face_trt_engines.py -- Build TensorRT FP16 engines for SCRFD and ArcFace.

Converts ONNX model files to optimized TensorRT engines with FP16 precision
for fast inference on Jetson Orin Nano Super.

Usage:
    python3 build_face_trt_engines.py \
        --scrfd-onnx /path/to/scrfd_2.5g_bnkps.onnx \
        --arcface-onnx /path/to/arcface_r50.onnx \
        --output-dir /mnt/nvme/saltybot/models \
        --fp16 --workspace-mb 1024
"""

import argparse
import os
import time


def build_engine(onnx_path: str, engine_path: str, fp16: bool, workspace_mb: int):
    """Build a TensorRT engine from an ONNX model.

    Args:
        onnx_path: Path to the source ONNX model file.
        engine_path: Output path for the serialized TensorRT engine.
        fp16: Enable FP16 precision.
        workspace_mb: Maximum workspace size in megabytes.
    """
    import tensorrt as trt

    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}')
    t0 = time.monotonic()

    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)}')
            raise RuntimeError(f'Failed to parse {onnx_path}')

    parse_time = time.monotonic() - t0
    print(f'  Parsed in {parse_time:.1f}s')

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE,
                                  workspace_mb * (1 << 20))
    if fp16:
        if builder.platform_has_fast_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
            print('  FP16 enabled.')
        else:
            print('  Warning: FP16 not supported on this platform, using FP32.')

    print(f'Building engine (this may take several minutes)...')
    t0 = time.monotonic()
    serialized = builder.build_serialized_network(network, config)
    build_time = time.monotonic() - t0

    if serialized is None:
        raise RuntimeError('Engine build failed.')

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

    size_mb = os.path.getsize(engine_path) / (1 << 20)
    print(f'  Engine saved: {engine_path} ({size_mb:.1f} MB, built in {build_time:.1f}s)')


def main():
    """Main entry point for TRT engine building."""
    parser = argparse.ArgumentParser(
        description='Build TensorRT FP16 engines for SCRFD and ArcFace.')
    parser.add_argument('--scrfd-onnx', type=str, default='',
                        help='Path to SCRFD ONNX model.')
    parser.add_argument('--arcface-onnx', type=str, default='',
                        help='Path to ArcFace ONNX model.')
    parser.add_argument('--output-dir', type=str,
                        default='/mnt/nvme/saltybot/models',
                        help='Output directory for engine files.')
    parser.add_argument('--fp16', action='store_true', default=True,
                        help='Enable FP16 precision (default: True).')
    parser.add_argument('--no-fp16', action='store_false', dest='fp16',
                        help='Disable FP16 (use FP32 only).')
    parser.add_argument('--workspace-mb', type=int, default=1024,
                        help='TRT workspace size in MB (default: 1024).')
    args = parser.parse_args()

    if not args.scrfd_onnx and not args.arcface_onnx:
        parser.error('At least one of --scrfd-onnx or --arcface-onnx is required.')

    if args.scrfd_onnx:
        engine_path = os.path.join(args.output_dir, 'scrfd_2.5g.engine')
        print(f'\n=== Building SCRFD engine ===')
        build_engine(args.scrfd_onnx, engine_path, args.fp16, args.workspace_mb)

    if args.arcface_onnx:
        engine_path = os.path.join(args.output_dir, 'arcface_r50.engine')
        print(f'\n=== Building ArcFace engine ===')
        build_engine(args.arcface_onnx, engine_path, args.fp16, args.workspace_mb)

    print('\nDone.')


if __name__ == '__main__':
    main()
