#!/usr/bin/env python3
"""Export a YOLO .pt model to TensorRT .engine.

Usage:
    python3 export_engine.py chicken-detection-model-v26n-300e-best-2026-05-02-NEW.pt
    python3 export_engine.py model.pt --imgsz 640 --half --workspace 4
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path


def export_engine(
    model_path: str | Path,
    *,
    imgsz: int = 640,
    half: bool = True,
    int8: bool = False,
    batch: int = 1,
    workspace: int = 4,     # GB
    simplify: bool = True,
    opset: int = 17,
    verbose: bool = True,
) -> str:
    from ultralytics import YOLO

    model = YOLO(model_path, task="detect")

    output = model.export(
        format="engine",
        imgsz=imgsz,
        half=half,
        int8=int8,
        batch=batch,
        workspace=workspace,
        simplify=simplify,
        opset=opset,
        verbose=verbose,
    )

    print(f"\nExported to: {output}")
    return str(output)


def main():
    parser = argparse.ArgumentParser(description="Export YOLO .pt → TensorRT .engine")
    parser.add_argument("model", help="Path to .pt model file")
    parser.add_argument("--imgsz", type=int, default=640, help="Input image size (default: 640)")
    parser.add_argument("--half", action="store_true", default=True, help="FP16 precision (default: on)")
    parser.add_argument("--no-half", dest="half", action="store_false", help="FP32 precision")
    parser.add_argument("--int8", action="store_true", help="INT8 quantization (needs calibration)")
    parser.add_argument("--batch", type=int, default=1, help="Batch size (default: 1)")
    parser.add_argument("--workspace", type=int, default=4, help="GPU workspace in GB (default: 4)")
    parser.add_argument("--opset", type=int, default=17, help="ONNX opset version (default: 17)")
    parser.add_argument("--quiet", action="store_true", help="Suppress verbose output")
    args = parser.parse_args()

    if not Path(args.model).exists():
        print(f"error: model file not found: {args.model}", file=sys.stderr)
        sys.exit(1)

    export_engine(
        args.model,
        imgsz=args.imgsz,
        half=args.half,
        int8=args.int8,
        batch=args.batch,
        workspace=args.workspace,
        opset=args.opset,
        verbose=not args.quiet,
    )


if __name__ == "__main__":
    main()