"""Run the trained ball detector on one frame and save an annotated image."""

from __future__ import annotations

import argparse

from pathlib import Path

import cv2

from rich.console import Console

ROOT = Path(__file__).resolve().parents[1]

MODEL = ROOT / "configs" / "models" / "mlb26_ball.pt"

console = Console()

def main() -> int:

    ap = argparse.ArgumentParser(description="Infer MLB ball model on one image.")

    ap.add_argument("image", type=Path)

    ap.add_argument("--model", type=Path, default=MODEL)

    ap.add_argument("--conf", type=float, default=0.25)

    ap.add_argument("--out", type=Path, default=ROOT / "logs" / "yolo_infer.png")

    args = ap.parse_args()

    from ultralytics import YOLO

    model = YOLO(str(args.model))

    results = model.predict(source=str(args.image), conf=args.conf, imgsz=640, verbose=False)

    annotated = results[0].plot()

    args.out.parent.mkdir(parents=True, exist_ok=True)

    cv2.imwrite(str(args.out), annotated)

    console.print(f"[green]Saved annotated output -> {args.out}[/green]")

    for box in results[0].boxes:

        xyxy = [float(x) for x in box.xyxy[0]]

        conf = float(box.conf[0])

        console.print(f"ball conf={conf:.3f} xyxy={xyxy}")

    return 0

if __name__ == "__main__":

    raise SystemExit(main())