"""Tiny OpenCV YOLO labeler for class 0 = ball.

Controls:

  drag left mouse: draw ball box, save, and go next

  s, Space, or Enter: save label and go next

  n: next without saving/changing

  b: previous

  c: clear current label

  e: save empty/no-ball label and go next

  q or Esc: quit

"""

from __future__ import annotations

import argparse

from pathlib import Path

import cv2

from rich.console import Console

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

DATA = ROOT / "datasets" / "mlb26_ball_yolo"

IMG_DIR = DATA / "images" / "all"

LBL_DIR = DATA / "labels" / "all"

console = Console()

class LabelState:

    def __init__(self, image, scale: float):

        self.image = image

        self.scale = scale

        self.box = None

        self.drag_start = None

        self.drag_now = None

        self.auto_advance = False

    def set_box_from_display(self, x1, y1, x2, y2):

        ox1 = int(min(x1, x2) / self.scale)

        oy1 = int(min(y1, y2) / self.scale)

        ox2 = int(max(x1, x2) / self.scale)

        oy2 = int(max(y1, y2) / self.scale)

        h, w = self.image.shape[:2]

        ox1, ox2 = max(0, ox1), min(w - 1, ox2)

        oy1, oy2 = max(0, oy1), min(h - 1, oy2)

        if ox2 - ox1 >= 2 and oy2 - oy1 >= 2:

            self.box = (ox1, oy1, ox2, oy2)

def yolo_line(box, w, h) -> str:

    x1, y1, x2, y2 = box

    xc = ((x1 + x2) / 2.0) / w

    yc = ((y1 + y2) / 2.0) / h

    bw = (x2 - x1) / w

    bh = (y2 - y1) / h

    return f"0 {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}\n"

def load_label(path: Path, w: int, h: int):

    if not path.exists() or not path.read_text(encoding="utf-8").strip():

        return None

    parts = path.read_text(encoding="utf-8").split()

    if len(parts) < 5:

        return None

    _, xc, yc, bw, bh = parts[:5]

    xc, yc, bw, bh = map(float, (xc, yc, bw, bh))

    x1 = int((xc - bw / 2.0) * w)

    y1 = int((yc - bh / 2.0) * h)

    x2 = int((xc + bw / 2.0) * w)

    y2 = int((yc + bh / 2.0) * h)

    return x1, y1, x2, y2

def draw(state: LabelState, title: str):

    img = state.image.copy()

    if state.box:

        x1, y1, x2, y2 = state.box

        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 255), 2)

    disp = cv2.resize(img, None, fx=state.scale, fy=state.scale, interpolation=cv2.INTER_AREA)

    if state.drag_start and state.drag_now:

        cv2.rectangle(disp, state.drag_start, state.drag_now, (0, 255, 0), 2)

    cv2.putText(disp, title, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 4, cv2.LINE_AA)

    cv2.putText(disp, title, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2, cv2.LINE_AA)

    help_text = "Draw tight box around baseball. Mouse-up saves + next. b=back c=clear e=empty q=quit"

    cv2.rectangle(disp, (8, disp.shape[0] - 42), (min(disp.shape[1] - 8, 980), disp.shape[0] - 8), (0, 0, 0), -1)

    cv2.putText(disp, help_text, (16, disp.shape[0] - 18), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 2, cv2.LINE_AA)

    return disp

def main() -> int:

    ap = argparse.ArgumentParser(description="Label YOLO ball boxes.")

    ap.add_argument("--images", type=Path, default=IMG_DIR)

    ap.add_argument("--labels", type=Path, default=LBL_DIR)

    ap.add_argument("--scale", type=float, default=0.65)

    ap.add_argument("--from-start", action="store_true", help="start at first image instead of first unlabeled image")

    ap.add_argument("--start-index", type=int, default=0, help="1-based index to start from")

    ap.add_argument("--start-name", default="", help="exact image filename to start from")

    ap.add_argument("--manual", action="store_true", help="do not auto-save/advance on mouse release")

    args = ap.parse_args()

    args.labels.mkdir(parents=True, exist_ok=True)

    images = sorted([p for p in args.images.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png"}])

    if not images:

        console.print(f"[red]No images found in {args.images}[/red]")

        return 2

    idx = 0

    if args.start_name:

        found = next((i for i, p in enumerate(images) if p.name == args.start_name), None)

        if found is None:

            console.print(f"[red]--start-name not found: {args.start_name}[/red]")

            return 2

        idx = found

    elif args.start_index > 0:

        idx = min(len(images) - 1, max(0, args.start_index - 1))

    elif not args.from_start:

        for i, image_path in enumerate(images):

            if not (args.labels / f"{image_path.stem}.txt").exists():

                idx = i

                break

    win = "label ball"

    while 0 <= idx < len(images):

        path = images[idx]

        img = cv2.imread(str(path))

        if img is None:

            idx += 1

            continue

        h, w = img.shape[:2]

        label_path = args.labels / f"{path.stem}.txt"

        state = LabelState(img, args.scale)

        state.box = load_label(label_path, w, h)

        def on_mouse(event, x, y, flags, param):

            if event == cv2.EVENT_LBUTTONDOWN:

                state.drag_start = (x, y)

                state.drag_now = (x, y)

            elif event == cv2.EVENT_MOUSEMOVE and state.drag_start:

                state.drag_now = (x, y)

            elif event == cv2.EVENT_LBUTTONUP and state.drag_start:

                state.set_box_from_display(state.drag_start[0], state.drag_start[1], x, y)

                state.drag_start = None

                state.drag_now = None

                if state.box and not args.manual:

                    state.auto_advance = True

        cv2.namedWindow(win, cv2.WINDOW_NORMAL)

        cv2.setMouseCallback(win, on_mouse)

        while True:

            title = f"{idx + 1}/{len(images)} {path.name}  auto-save on mouse-up | s save | b back | c clear | q quit"

            cv2.imshow(win, draw(state, title))

            key = cv2.waitKey(16) & 0xFF

            if state.auto_advance:

                label_path.write_text(yolo_line(state.box, w, h), encoding="utf-8")

                idx += 1

                break

            if key in (ord("q"), 27):

                cv2.destroyAllWindows()

                return 0

            if key == ord("c"):

                state.box = None

                if label_path.exists():

                    label_path.unlink()

            if key in (ord("s"), 13, 32):

                if state.box:

                    label_path.write_text(yolo_line(state.box, w, h), encoding="utf-8")

                idx += 1

                break

            if key == ord("e"):

                label_path.write_text("", encoding="utf-8")

                idx += 1

                break

            if key == ord("n"):

                idx += 1

                break

            if key == ord("b"):

                idx = max(0, idx - 1)

                break

    cv2.destroyAllWindows()

    console.print("[green]Labeling complete.[/green]")

    return 0

if __name__ == "__main__":

    raise SystemExit(main())