"""Visual overlay for live debugging.

Subscribes to:

  - capture frames         (capture.publish_endpoint)

  - ball/pitch events      (cv.ball_events_endpoint)

  - PCI events             (cv.pci_events_endpoint)

Draws the most recent detections on each frame. Press q or ESC to quit.

Run capture/ingest.py, cv/ball_tracker.py, cv/pci_tracker.py first.

"""

from __future__ import annotations

import argparse

import json

import sys

import time

from pathlib import Path

import cv2

import numpy as np

import zmq

from rich.console import Console

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

from cv._common import (

    event_subscriber,

    iter_latest_frames,

    load_config,

    make_frame_subscriber,

)

console = Console()

COLOR_BALL = (0, 255, 255)

COLOR_BALL_PRED = (0, 128, 255)

COLOR_PCI = (0, 255, 0)

COLOR_PLATE = (255, 255, 255)

COLOR_BG = (0, 0, 0)

def drain_latest_json(sock: zmq.Socket) -> list[dict]:

    """Pull every pending JSON event without blocking."""

    events: list[dict] = []

    while True:

        try:

            raw = sock.recv(flags=zmq.NOBLOCK)

        except zmq.Again:

            break

        try:

            events.append(json.loads(raw.decode("utf-8")))

        except Exception:

            pass

    return events

def draw_hud(frame: np.ndarray, lines: list[str]) -> None:

    y = 28

    for line in lines:

        cv2.putText(frame, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 3, cv2.LINE_AA)

        cv2.putText(frame, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)

        y += 22

def main() -> int:

    ap = argparse.ArgumentParser(description="pitch-tracker-cv live viewer.")

    ap.add_argument("--scale", type=float, default=0.75, help="Display scale factor (0.5 = half-size).")

    ap.add_argument("--no-gui", action="store_true", help="Pull events and print stats without opening a window.")

    args = ap.parse_args()

    cfg = load_config()

    cap_ep = cfg["capture"]["publish_endpoint"]

    ball_ep = cfg["cv"]["ball_events_endpoint"]

    pci_ep = cfg["cv"]["pci_events_endpoint"]

    plate_y_frac = float(cfg["cv"].get("plate_y_frac", 0.72))

    ctx = zmq.Context.instance()

    frame_sub = make_frame_subscriber(cap_ep, ctx=ctx)

    ball_sub = event_subscriber(ball_ep, ctx=ctx)

    pci_sub = event_subscriber(pci_ep, ctx=ctx)

    console.print(f"[green]viewer[/green] cap={cap_ep} ball={ball_ep} pci={pci_ep}")

    last_ball: dict | None = None

    last_pred: dict | None = None

    last_pci: dict | None = None

    frames = 0

    t_start = time.perf_counter()

    win = "pitch-tracker-cv viewer"

    if not args.no_gui:

        cv2.namedWindow(win, cv2.WINDOW_NORMAL)

    try:

        for meta, frame in iter_latest_frames(frame_sub, timeout_ms=3000):

            frames += 1

            for ev in drain_latest_json(ball_sub):

                if ev.get("type") == "ball_track":

                    last_ball = ev

                elif ev.get("type") == "pitch_pred":

                    last_pred = ev

                elif ev.get("type") == "ball_miss":

                    if last_ball and ev["ts_ns"] - last_ball.get("ts_ns", 0) > 3e8:

                        last_ball = None

            for ev in drain_latest_json(pci_sub):

                if ev.get("type") == "pci_track":

                    last_pci = ev

                elif ev.get("type") == "pci_miss":

                    if last_pci and ev["ts_ns"] - last_pci.get("ts_ns", 0) > 3e8:

                        last_pci = None

            if args.no_gui:

                if frames % 60 == 0:

                    console.print(

                        f"[dim]  viewer: {frames} frames  "

                        f"ball={'y' if last_ball else 'n'}  "

                        f"pred={'y' if last_pred else 'n'}  "

                        f"pci={'y' if last_pci else 'n'}[/dim]"

                    )

                continue

            h, w = frame.shape[:2]

            overlay = frame.copy()

            plate_y = int(h * plate_y_frac)

            cv2.line(overlay, (0, plate_y), (w, plate_y), COLOR_PLATE, 1, cv2.LINE_AA)

            if last_ball is not None:

                bx, by, br = int(last_ball["x"]), int(last_ball["y"]), int(last_ball["r"])

                cv2.circle(overlay, (bx, by), max(br, 4), COLOR_BALL, 2)

                cv2.circle(overlay, (bx, by), 2, COLOR_BALL, -1)

            if last_pred is not None:

                px = int(last_pred["plate_x"])

                py = int(last_pred["plate_y"])

                cv2.drawMarker(overlay, (px, py), COLOR_BALL_PRED, cv2.MARKER_CROSS, 28, 2)

                cv2.putText(

                    overlay,

                    f"eta {last_pred['eta_ms']:.0f}ms",

                    (px + 14, py - 8),

                    cv2.FONT_HERSHEY_SIMPLEX,

                    0.6,

                    COLOR_BALL_PRED,

                    2,

                    cv2.LINE_AA,

                )

            if last_pci is not None:

                pcx, pcy, pcr = int(last_pci["x"]), int(last_pci["y"]), int(last_pci["r"])

                cv2.circle(overlay, (pcx, pcy), pcr, COLOR_PCI, 2)

                cv2.circle(overlay, (pcx, pcy), 3, COLOR_PCI, -1)

            dt = time.perf_counter() - t_start

            fps = frames / dt if dt > 0 else 0.0

            draw_hud(overlay, [

                f"frames {frames}  display {fps:.1f} fps  seq {meta['seq']}",

                f"ball {'Y' if last_ball else '-'}   pred {'Y' if last_pred else '-'}   pci {'Y' if last_pci else '-'}",

            ])

            if args.scale != 1.0:

                overlay = cv2.resize(overlay, None, fx=args.scale, fy=args.scale, interpolation=cv2.INTER_AREA)

            cv2.imshow(win, overlay)

            key = cv2.waitKey(1) & 0xFF

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

                break

    except TimeoutError as e:

        console.print(f"[red]viewer: {e}. Is capture/ingest.py running?[/red]")

        return 2

    except KeyboardInterrupt:

        console.print("[yellow]viewer interrupted.[/yellow]")

    finally:

        frame_sub.close(0)

        ball_sub.close(0)

        pci_sub.close(0)

        if not args.no_gui:

            cv2.destroyAllWindows()

    return 0

if __name__ == "__main__":

    sys.exit(main())