"""Export the trained model to edge-deployment formats (plan 3.2).

Requires the `export` dependency group. From repo root:

    uv run --group export python scripts/export_model.py                      # onnx + openvino + coreml
    uv run --group export python scripts/export_model.py --formats onnx
    uv run --group export python scripts/export_model.py --formats engine     # TensorRT — NVIDIA GPU only (Colab)

Artifacts land next to the weights (runs/detect/.../weights/) and are
git-ignored — re-run this script to regenerate them. Benchmark them against the
PyTorch baseline with scripts/bench_exports.py.
"""

from __future__ import annotations

import argparse
from pathlib import Path

from ultralytics import YOLO

WEIGHTS = "runs/detect/yolov8n_v1_train/weights/best.pt"
DEFAULT_FORMATS = ["onnx", "openvino", "coreml"]


def dir_size_mb(path: Path) -> float:
    if path.is_file():
        return path.stat().st_size / 1e6
    return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1e6


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--weights", default=WEIGHTS)
    p.add_argument(
        "--formats",
        default=",".join(DEFAULT_FORMATS),
        help="comma-separated Ultralytics export formats (onnx, openvino, coreml, engine, tflite, ...)",
    )
    p.add_argument("--imgsz", type=int, default=640)
    p.add_argument("--half", action="store_true", help="FP16 export (GPU formats like engine)")
    args = p.parse_args()

    exported: list[tuple[str, Path]] = []
    for fmt in args.formats.split(","):
        fmt = fmt.strip()
        # Fresh model per format: export() mutates model state in some backends.
        model = YOLO(args.weights)
        print(f"\n=== Exporting {fmt} ===")
        out = model.export(format=fmt, imgsz=args.imgsz, half=args.half)
        exported.append((fmt, Path(out)))

    print(f"\n{'Format':<10} {'Size (MB)':>10}  Path")
    print("-" * 60)
    for fmt, path in exported:
        print(f"{fmt:<10} {dir_size_mb(path):>10.1f}  {path}")


if __name__ == "__main__":
    main()