"""Capture live frames for one-class MLB baseball YOLO training.

Saves raw frames to datasets/mlb26_ball_yolo/images/all. Label these images

with class 0 = ball, then run yolo_split_dataset.py.

"""

from __future__ import annotations

import argparse

import shutil

import sys

import time

from pathlib import Path

import cv2

from rich.console import Console

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from capture.ingest import open_capture

from cv._common import load_config

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

OUT_DIR = ROOT / "datasets" / "mlb26_ball_yolo" / "images" / "all"

console = Console()

def main() -> int:

    ap = argparse.ArgumentParser(description="Capture frames for YOLO ball labeling.")

    ap.add_argument("--duration", type=float, default=120.0, help="Capture duration in seconds.")

    ap.add_argument("--fps", type=float, default=12.0, help="Saved frame rate.")

    ap.add_argument("--clear", action="store_true", help="Clear images/all before capture.")

    ap.add_argument("--prefix", default="mlb26", help="Output filename prefix.")

    args = ap.parse_args()

    if args.clear and OUT_DIR.exists():

        shutil.rmtree(OUT_DIR)

    OUT_DIR.mkdir(parents=True, exist_ok=True)

    cap = open_capture(load_config())

    if cap is None or not cap.isOpened():

        console.print("[red]Could not open capture card.[/red]")

        return 2

    interval = 1.0 / max(args.fps, 0.1)

    deadline = time.perf_counter() + args.duration

    next_save = time.perf_counter()

    saved = 0

    read_count = 0

    console.print(f"[cyan]Capturing {args.duration:.0f}s at {args.fps:.1f} saved fps -> {OUT_DIR}[/cyan]")

    try:

        while time.perf_counter() < deadline:

            ok, frame = cap.read()

            if not ok or frame is None:

                continue

            read_count += 1

            now = time.perf_counter()

            if now < next_save:

                continue

            out = OUT_DIR / f"{args.prefix}_{int(time.time() * 1000)}_{saved:05d}.jpg"

            cv2.imwrite(str(out), frame, [cv2.IMWRITE_JPEG_QUALITY, 95])

            saved += 1

            next_save = now + interval

            if saved % 50 == 0:

                console.print(f"  saved {saved} frames")

    finally:

        cap.release()

    console.print(f"[green]Saved {saved} frames. Capture reads: {read_count}.[/green]")

    return 0

if __name__ == "__main__":

    raise SystemExit(main())