#!/usr/bin/env python3
"""TensorRT engine conversion script.

Converts ONNX models to TensorRT FP16 .engine files for Jetson deployment.
Wraps trtexec CLI (available on Jetson with JetPack installed).

Usage:
    # Convert a single model
    python scripts/convert_to_trt.py --onnx weights/litesam.onnx --output /opt/engines/litesam.engine

    # Convert all known models from weights/ to engines/
    python scripts/convert_to_trt.py --all --onnx-dir weights/ --engine-dir /opt/engines/

    # Dry run (check trtexec availability)
    python scripts/convert_to_trt.py --check

Requires:
    - NVIDIA TensorRT (trtexec in PATH) — available on Jetson with JetPack
    - NOT available on dev/CI machines — script exits cleanly with a message
"""

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys

# Known models and their expected ONNX filenames
_MODELS = {
    "superpoint": "superpoint.onnx",
    "lightglue": "lightglue.onnx",
    "xfeat": "xfeat.onnx",
    "dinov2": "dinov2.onnx",
    "litesam": "litesam.onnx",
}


def find_trtexec() -> str | None:
    """Find trtexec binary in PATH or common Jetson locations."""
    path = shutil.which("trtexec")
    if path:
        return path
    # Common Jetson paths
    for candidate in [
        "/usr/src/tensorrt/bin/trtexec",
        "/usr/local/cuda/bin/trtexec",
    ]:
        if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
            return candidate
    return None


def convert_onnx_to_engine(
    onnx_path: str,
    engine_path: str,
    fp16: bool = True,
    workspace_mb: int = 1024,
    trtexec_path: str | None = None,
) -> bool:
    """Run trtexec to convert ONNX → TensorRT engine.

    Returns True on success, False on failure.
    """
    trtexec = trtexec_path or find_trtexec()
    if not trtexec:
        print("ERROR: trtexec not found. Install TensorRT or run on Jetson with JetPack.")
        return False

    if not os.path.isfile(onnx_path):
        print(f"ERROR: ONNX file not found: {onnx_path}")
        return False

    os.makedirs(os.path.dirname(engine_path) or ".", exist_ok=True)

    cmd = [
        trtexec,
        f"--onnx={onnx_path}",
        f"--saveEngine={engine_path}",
        f"--workspace={workspace_mb}",
    ]
    if fp16:
        cmd.append("--fp16")

    print(f"  Converting: {onnx_path} → {engine_path}")
    print(f"  Command: {' '.join(cmd)}")

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        if result.returncode == 0:
            size_mb = os.path.getsize(engine_path) / (1024 * 1024)
            print(f"  OK: {engine_path} ({size_mb:.1f} MB)")
            return True
        else:
            print(f"  FAIL (exit {result.returncode})")
            if result.stderr:
                print(f"  stderr: {result.stderr[:500]}")
            return False
    except subprocess.TimeoutExpired:
        print("  FAIL: trtexec timed out (>600s)")
        return False
    except FileNotFoundError:
        print(f"  FAIL: trtexec not found at {trtexec}")
        return False


def main() -> int:
    parser = argparse.ArgumentParser(description="Convert ONNX models to TensorRT FP16 engines")
    parser.add_argument("--onnx", help="Input ONNX model path")
    parser.add_argument("--output", help="Output .engine path")
    parser.add_argument("--all", action="store_true", help="Convert all known models")
    parser.add_argument("--onnx-dir", default="weights", help="Directory with ONNX files")
    parser.add_argument("--engine-dir", default="/opt/engines", help="Output engine directory")
    parser.add_argument("--no-fp16", action="store_true", help="Disable FP16 (use FP32)")
    parser.add_argument("--workspace", type=int, default=1024, help="TRT workspace (MB)")
    parser.add_argument("--check", action="store_true", help="Check trtexec availability only")
    args = parser.parse_args()

    # Check mode
    if args.check:
        trtexec = find_trtexec()
        if trtexec:
            print(f"trtexec found: {trtexec}")
            return 0
        else:
            print("trtexec not found. Not on Jetson or TensorRT not installed.")
            return 1

    fp16 = not args.no_fp16

    # Single model conversion
    if args.onnx and args.output:
        ok = convert_onnx_to_engine(args.onnx, args.output, fp16=fp16, workspace_mb=args.workspace)
        return 0 if ok else 1

    # Batch conversion
    if args.all:
        trtexec = find_trtexec()
        if not trtexec:
            print("trtexec not found. Conversion requires Jetson with JetPack installed.")
            return 1

        print(f"Converting all known models from {args.onnx_dir}/ → {args.engine_dir}/")
        success = 0
        fail = 0
        for model_name, onnx_file in _MODELS.items():
            onnx_path = os.path.join(args.onnx_dir, onnx_file)
            engine_path = os.path.join(args.engine_dir, f"{model_name}.engine")
            if not os.path.isfile(onnx_path):
                print(f"  SKIP {model_name}: {onnx_path} not found")
                continue
            ok = convert_onnx_to_engine(
                onnx_path, engine_path,
                fp16=fp16, workspace_mb=args.workspace, trtexec_path=trtexec,
            )
            if ok:
                success += 1
            else:
                fail += 1

        print(f"\nDone: {success} converted, {fail} failed")
        return 1 if fail > 0 else 0

    parser.print_help()
    return 1


if __name__ == "__main__":
    sys.exit(main())