#!/usr/bin/env python3
"""
Build YOLOv8n TensorRT FP16 engine for Orin Nano.
"""

import argparse
import sys
from pathlib import Path


def build_engine(output_path: Path, workspace_mb: int = 2048) -> bool:
    """Download YOLOv8n, export ONNX, and convert to TensorRT."""
    try:
        from ultralytics import YOLO
    except ImportError:
        print("ERROR: ultralytics not installed. Install with: pip install ultralytics")
        return False

    try:
        import tensorrt as trt
    except ImportError:
        print("ERROR: TensorRT not installed")
        return False

    print("[*] Loading YOLOv8n from Ultralytics...")
    model = YOLO("yolov8n.pt")

    print("[*] Exporting to ONNX...")
    onnx_path = output_path.parent / "yolov8n.onnx"
    model.export(format="onnx", opset=12)

    # Move ONNX to desired location
    import shutil
    onnx_src = Path("yolov8n.onnx")
    if onnx_src.exists():
        shutil.move(str(onnx_src), str(onnx_path))
        print(f"[+] ONNX exported to {onnx_path}")
    else:
        print(f"ERROR: ONNX export not found")
        return False

    print("[*] Converting ONNX to TensorRT FP16...")
    try:
        import polygraphy
        from polygraphy.backend.trt import engine_from_network, save_engine
        from polygraphy.backend.onnx import BytesFromOnnx
    except ImportError:
        print("ERROR: polygraphy not installed. Install with: pip install polygraphy")
        return False

    try:
        # Build TRT engine
        logger = trt.Logger(trt.Logger.INFO)
        onnx_bytes = BytesFromOnnx(str(onnx_path))()

        engine = engine_from_network(
            onnx_bytes,
            config_kwargs={
                "flags": [trt.BuilderFlag.FP16],
                "max_workspace_size": workspace_mb * 1024 * 1024,
            },
            logger=logger,
        )()

        # Save engine
        save_engine(engine, str(output_path))
        print(f"[+] TensorRT engine saved to {output_path}")
        return True

    except Exception as e:
        print(f"ERROR: Failed to convert to TensorRT: {e}")
        return False


def benchmark_engine(engine_path: Path, num_iterations: int = 100) -> None:
    """Benchmark TensorRT engine latency."""
    try:
        import tensorrt as trt
        import pycuda.driver as cuda
        import pycuda.autoinit
        import numpy as np
        import time
    except ImportError as e:
        print(f"ERROR: Missing dependency: {e}")
        return

    if not engine_path.exists():
        print(f"ERROR: Engine not found: {engine_path}")
        return

    print(f"\n[*] Benchmarking {engine_path} ({num_iterations} iterations)...")

    try:
        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f:
            engine = trt.Runtime(logger).deserialize_cuda_engine(f.read())

        context = engine.create_execution_context()
        stream = cuda.Stream()

        # Prepare input (1, 3, 640, 640)
        h_input = cuda.pagelocked_empty(1 * 3 * 640 * 640, np.float32)
        d_input = cuda.mem_alloc(h_input.nbytes)

        # Prepare output (1, 84, 8400)
        h_output = cuda.pagelocked_empty(1 * 84 * 8400, np.float32)
        d_output = cuda.mem_alloc(h_output.nbytes)

        bindings = [int(d_input), int(d_output)]

        # Warmup
        for _ in range(10):
            cuda.memcpy_htod_async(d_input, h_input, stream)
            context.execute_async_v2(bindings, stream.handle)
            cuda.memcpy_dtoh_async(h_output, d_output, stream)
            stream.synchronize()

        # Benchmark
        times = []
        for _ in range(num_iterations):
            cuda.memcpy_htod_async(d_input, h_input, stream)
            start = time.time()
            context.execute_async_v2(bindings, stream.handle)
            cuda.memcpy_dtoh_async(h_output, d_output, stream)
            stream.synchronize()
            elapsed = time.time() - start
            times.append(elapsed * 1000)  # ms

        mean_latency = np.mean(times)
        std_latency = np.std(times)
        min_latency = np.min(times)
        max_latency = np.max(times)
        throughput = 1000.0 / mean_latency

        print(f"[+] Latency: {mean_latency:.2f}ms ± {std_latency:.2f}ms (min={min_latency:.2f}ms, max={max_latency:.2f}ms)")
        print(f"[+] Throughput: {throughput:.1f} FPS")

    except Exception as e:
        print(f"ERROR: Benchmark failed: {e}")
    finally:
        d_input.free()
        d_output.free()


def main():
    parser = argparse.ArgumentParser(description="Build YOLOv8n TensorRT engine")
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("/mnt/nvme/saltybot/models/yolov8n.engine"),
        help="Output engine path"
    )
    parser.add_argument(
        "--workspace",
        type=int,
        default=2048,
        help="TensorRT workspace size in MB"
    )
    parser.add_argument(
        "--benchmark",
        action="store_true",
        help="Benchmark the engine after building"
    )

    args = parser.parse_args()

    # Create output directory
    args.output.parent.mkdir(parents=True, exist_ok=True)

    print(f"[*] Building YOLOv8n TensorRT engine")
    print(f"[*] Output: {args.output}")
    print(f"[*] Workspace: {args.workspace} MB")

    success = build_engine(args.output, args.workspace)

    if success and args.benchmark:
        benchmark_engine(args.output)

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
