"""Capture card ingest service.

Smoke mode (--smoke): 10s test that prints FPS + latency and saves one frame.

Default: subscribe-side-agnostic ZMQ PUB of JPEG frames on capture.publish_endpoint.

"""

from __future__ import annotations

import argparse

import json

import sys

import time

from pathlib import Path

import cv2

import numpy as np

import yaml

import zmq

from rich.console import Console

console = Console()

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

CONFIG_PATH = ROOT / "configs" / "runtime.yaml"

LOG_DIR = ROOT / "logs"

def load_config() -> dict:

    with open(CONFIG_PATH, "r", encoding="utf-8") as f:

        return yaml.safe_load(f)

def probe_devices(max_index: int = 5, backend: int = cv2.CAP_DSHOW) -> list[tuple[int, int, int]]:

    found: list[tuple[int, int, int]] = []

    for i in range(max_index + 1):

        cap = cv2.VideoCapture(i, backend)

        if cap.isOpened():

            w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

            h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

            found.append((i, w, h))

        cap.release()

    return found

def _try_open(idx: int, backend: int, w: int, h: int, fps: int) -> cv2.VideoCapture | None:

    cap = cv2.VideoCapture(idx, backend)

    if not cap.isOpened():

        cap.release()

        return None

    cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"MJPG"))

    cap.set(cv2.CAP_PROP_FRAME_WIDTH, w)

    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, h)

    cap.set(cv2.CAP_PROP_FPS, fps)

    ok, frame = cap.read()

    if not ok or frame is None:

        cap.release()

        return None

    return cap

def open_capture(cfg: dict) -> cv2.VideoCapture | None:

    cap_cfg = cfg["capture"]

    idx = cap_cfg.get("device_index", 0)

    w = cap_cfg.get("width", 1920)

    h = cap_cfg.get("height", 1080)

    fps = cap_cfg.get("fps", 60)

    candidate_indices = [idx]

    for extra in range(0, 6):

        if extra not in candidate_indices:

            candidate_indices.append(extra)

    backends = [cv2.CAP_DSHOW, cv2.CAP_MSMF, cv2.CAP_ANY]

    for backend in backends:

        for cand_idx in candidate_indices:

            cap = _try_open(cand_idx, backend, w, h, fps)

            if cap is not None:

                console.print(

                    f"[dim]open_capture: index={cand_idx} backend={backend} "

                    f"size={int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))}x{int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))}[/dim]"

                )

                return cap

    return None

def smoke_test(cfg: dict, duration_s: float = 10.0) -> int:

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

    console.print("[bold cyan]Probing video devices 0..5 (DirectShow)...[/bold cyan]")

    devs = probe_devices()

    if not devs:

        console.print(

            "[bold red]No DirectShow video devices could be opened.[/bold red]\n"

            "  Check: capture card USB cable seated, card has power (LED on),\n"

            "  Xbox HDMI plugged into card HDMI IN (not OUT)."

        )

        return 2

    for i, w, h in devs:

        console.print(f"  device {i}: default size {w}x{h}")

    cap = open_capture(cfg)

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

        console.print(

            f"[bold red]Could not open configured device_index={cfg['capture']['device_index']}.[/bold red]\n"

            "  Try a different index from the probe results above and update configs/runtime.yaml."

        )

        return 2

    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))

    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    reported_fps = cap.get(cv2.CAP_PROP_FPS)

    console.print(

        f"[green]Opened device_index={cfg['capture']['device_index']} "

        f"at {w}x{h}, reported fps={reported_fps:.1f}[/green]"

    )

    console.print(f"[bold]Running {duration_s:.1f}s smoke test...[/bold]")

    frame_count = 0

    first_frame: np.ndarray | None = None

    last_frame: np.ndarray | None = None

    latencies_ms: list[float] = []

    t_start = time.perf_counter()

    deadline = t_start + duration_s

    while time.perf_counter() < deadline:

        t0 = time.perf_counter()

        ret, frame = cap.read()

        t1 = time.perf_counter()

        if not ret or frame is None:

            continue

        latencies_ms.append((t1 - t0) * 1000.0)

        if first_frame is None:

            first_frame = frame.copy()

        last_frame = frame

        frame_count += 1

    elapsed = time.perf_counter() - t_start

    cap.release()

    if frame_count == 0 or first_frame is None:

        console.print("[bold red]Zero frames captured in 10 seconds.[/bold red]")

        return 3

    fps = frame_count / elapsed if elapsed > 0 else 0.0

    lat = np.asarray(latencies_ms)

    console.print("[bold green]Smoke test results[/bold green]")

    console.print(f"  frames:        {frame_count}")

    console.print(f"  elapsed:       {elapsed:.2f}s")

    console.print(f"  measured FPS:  {fps:.2f}")

    console.print(f"  frame size:    {first_frame.shape[1]}x{first_frame.shape[0]}  (channels={first_frame.shape[2]})")

    console.print(

        f"  read latency:  mean {lat.mean():.2f}ms  "

        f"p50 {np.percentile(lat, 50):.2f}ms  "

        f"p95 {np.percentile(lat, 95):.2f}ms  "

        f"max {lat.max():.2f}ms"

    )

    out_first = LOG_DIR / "smoke_frame.png"

    cv2.imwrite(str(out_first), first_frame)

    console.print(f"[green]Saved first frame -> {out_first}[/green]")

    if last_frame is not None and frame_count > 1:

        out_last = LOG_DIR / "smoke_frame_last.png"

        cv2.imwrite(str(out_last), last_frame)

        console.print(f"[green]Saved last frame  -> {out_last}[/green]")

    return 0

def run_publisher(cfg: dict) -> int:

    cap = open_capture(cfg)

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

        console.print("[red]Capture device not available; publisher exiting.[/red]")

        return 2

    endpoint = cfg["capture"]["publish_endpoint"]

    jpeg_q = int(cfg["capture"].get("jpeg_quality", 80))

    ctx = zmq.Context.instance()

    sock = ctx.socket(zmq.PUB)

    sock.setsockopt(zmq.SNDHWM, 2)

    sock.bind(endpoint)

    console.print(f"[green]Publishing multipart [meta, jpeg] on {endpoint}[/green]  (Ctrl-C to stop)")

    seq = 0

    t_last_report = time.perf_counter()

    frames_since_report = 0

    try:

        while True:

            ret, frame = cap.read()

            if not ret or frame is None:

                continue

            ok, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, jpeg_q])

            if not ok:

                continue

            meta = {

                "seq": seq,

                "ts_ns": time.time_ns(),

                "h": int(frame.shape[0]),

                "w": int(frame.shape[1]),

            }

            sock.send_multipart([json.dumps(meta).encode("utf-8"), buf.tobytes()])

            seq += 1

            frames_since_report += 1

            now = time.perf_counter()

            if now - t_last_report >= 5.0:

                fps = frames_since_report / (now - t_last_report)

                console.print(f"[dim]  publisher: {fps:.1f} fps, last seq={seq - 1}[/dim]")

                t_last_report = now

                frames_since_report = 0

    except KeyboardInterrupt:

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

    finally:

        cap.release()

        sock.close(0)

        ctx.term()

    return 0

def main() -> None:

    p = argparse.ArgumentParser(description="pitch-tracker-cv capture card ingest.")

    p.add_argument("--smoke", action="store_true", help="10s smoke test: print FPS/latency, save one frame, exit.")

    p.add_argument("--duration", type=float, default=10.0, help="Smoke test duration in seconds.")

    args = p.parse_args()

    cfg = load_config()

    if args.smoke:

        sys.exit(smoke_test(cfg, duration_s=args.duration))

    sys.exit(run_publisher(cfg))

if __name__ == "__main__":

    main()