orisys-sdk/src/orisys/capture/source.py

"""OpenCV-backed capture source with optional reconnect."""

from __future__ import annotations

import os
import sys
import time
from typing import Any, Optional, Protocol, runtime_checkable

import cv2
import numpy as np

from .config import (
    CaptureConfig,
    CaptureHealth,
    CaptureStatus,
    ReconnectPolicy,
    SourceType,
)


@runtime_checkable
class VideoCaptureLike(Protocol):
    """Minimal protocol so tests can inject fake captures."""

    def isOpened(self) -> bool: ...
    def read(self) -> tuple[bool, Any]: ...
    def release(self) -> None: ...
    def get(self, propId: int) -> float: ...
    def set(self, propId: int, value: float) -> bool: ...


def normalize_source(source: SourceType | None) -> SourceType:
    """Accept camera index, numeric string, V4L path, or filesystem path."""
    if source is None:
        return 0
    if isinstance(source, str):
        stripped = source.strip()
        if not stripped:
            return 0
        if stripped.isdigit():
            return int(stripped)
        # Only strip trailing separators (Windows path quirk); keep leading '/'.
        return stripped.rstrip("/\\") or stripped
    return source


def is_live_camera(source: SourceType) -> bool:
    """True for camera indices and Linux V4L2 device paths."""
    if isinstance(source, int):
        return True
    if isinstance(source, str) and source.startswith("/dev/video"):
        return True
    return False


def _backend_candidates(preferred: str = "auto") -> list[tuple[str, int]]:
    preferred = (preferred or "auto").lower()
    candidates: list[tuple[str, int]] = []

    def add(name: str, attr: str) -> None:
        if hasattr(cv2, attr):
            candidates.append((name, getattr(cv2, attr)))

    if preferred in ("dshow", "directshow") and sys.platform == "win32":
        add("DSHOW", "CAP_DSHOW")
    elif preferred in ("msmf",) and sys.platform == "win32":
        add("MSMF", "CAP_MSMF")
    elif preferred in ("v4l2", "v4l") and sys.platform.startswith("linux"):
        add("V4L2", "CAP_V4L2")
    elif preferred not in ("auto", "any", ""):
        # Unknown preference: fall through to platform defaults.
        pass

    if not candidates:
        if sys.platform == "win32":
            # Prefer DSHOW to avoid MSMF HW-transform stalls on many USB cameras.
            add("DSHOW", "CAP_DSHOW")
            add("MSMF", "CAP_MSMF")
        elif sys.platform.startswith("linux"):
            add("V4L2", "CAP_V4L2")
        if not candidates:
            candidates.append(("ANY", cv2.CAP_ANY))
    return candidates


def _fourcc_code(fourcc: Optional[str]) -> Optional[int]:
    if not fourcc:
        return None
    text = str(fourcc).strip()
    if len(text) != 4:
        return None
    return cv2.VideoWriter.fourcc(*text)


def _decode_fourcc(value: float) -> str:
    fourcc_int = int(value)
    if fourcc_int <= 0:
        return "N/A"
    bytes_ = [(fourcc_int >> (8 * i)) & 0xFF for i in range(4)]
    return "".join(chr(b) if 32 <= b <= 126 else "." for b in bytes_)


def open_opencv_capture(
    source: SourceType,
    *,
    backend: str = "auto",
    width: Optional[int] = None,
    height: Optional[int] = None,
    fps: Optional[float] = None,
    fourcc: Optional[str] = "MJPG",
    apply_stream_props: bool = True,
) -> tuple[VideoCaptureLike, str]:
    """Open a ``cv2.VideoCapture`` trying preferred backends in order.

    Returns ``(capture, backend_name)``.
    """
    # Avoid MSMF HW transform stalls on Windows camera capture.
    if sys.platform == "win32":
        os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0")

    source = normalize_source(source)
    live = is_live_camera(source)
    last_error = "Failed to open video source."

    backends = _backend_candidates(backend) if live else [("ANY", cv2.CAP_ANY)]
    for backend_name, backend_api in backends:
        try:
            if live and backend_api != cv2.CAP_ANY:
                cap = cv2.VideoCapture(source, backend_api)
            else:
                cap = cv2.VideoCapture(source)
        except Exception as exc:  # pragma: no cover - OpenCV quirks
            last_error = str(exc)
            continue
        if not cap.isOpened():
            cap.release()
            last_error = f"Backend {backend_name} failed to open {source!r}"
            continue
        if live and apply_stream_props:
            code = _fourcc_code(fourcc)
            if code is not None:
                cap.set(cv2.CAP_PROP_FOURCC, code)
            if fps is not None:
                cap.set(cv2.CAP_PROP_FPS, float(fps))
            if width is not None:
                cap.set(cv2.CAP_PROP_FRAME_WIDTH, float(width))
            if height is not None:
                cap.set(cv2.CAP_PROP_FRAME_HEIGHT, float(height))
        return cap, backend_name

    raise RuntimeError(last_error)


class CaptureSource:
    """Owns a video capture handle, health state, and reconnect logic.

    Optionally accepts an injected ``factory`` for tests::

        CaptureSource(config, factory=lambda cfg: (FakeCap(), "FAKE"))
    """

    def __init__(
        self,
        config: CaptureConfig,
        *,
        factory=None,
    ):
        self.config = config
        self._factory = factory or self._default_factory
        self._cap: Optional[VideoCaptureLike] = None
        self._backend_name = ""
        self._status = CaptureStatus(source=normalize_source(config.source))
        self._reconnect_attempt = 0
        self._lock = __import__("threading").Lock()

    @staticmethod
    def _default_factory(config: CaptureConfig) -> tuple[VideoCaptureLike, str]:
        return open_opencv_capture(
            config.source,
            backend=config.backend,
            width=config.width,
            height=config.height,
            fps=config.fps,
            fourcc=config.fourcc,
            apply_stream_props=is_live_camera(normalize_source(config.source)),
        )

    @property
    def status(self) -> CaptureStatus:
        with self._lock:
            return CaptureStatus(
                health=self._status.health,
                source=self._status.source,
                last_error=self._status.last_error,
                last_frame_time=self._status.last_frame_time,
                last_seq=self._status.last_seq,
                dropped_frames=self._status.dropped_frames,
                reconnect_attempts=self._status.reconnect_attempts,
                consecutive_failures=self._status.consecutive_failures,
                width=self._status.width,
                height=self._status.height,
                fps=self._status.fps,
                backend_name=self._status.backend_name,
                fourcc=self._status.fourcc,
            )

    @property
    def cap(self) -> Optional[VideoCaptureLike]:
        """Raw capture handle (compat with Sensor.cap)."""
        return self._cap

    def is_open(self) -> bool:
        cap = self._cap
        return cap is not None and cap.isOpened()

    def open(self) -> None:
        with self._lock:
            self._status.health = CaptureHealth.CONNECTING
            self._status.last_error = None
        try:
            cap, backend_name = self._factory(self.config)
        except Exception as exc:
            with self._lock:
                self._status.health = CaptureHealth.FAILED
                self._status.last_error = str(exc)
            raise
        with self._lock:
            self._release_unlocked()
            self._cap = cap
            self._backend_name = backend_name
            self._refresh_stream_meta_unlocked()
            self._status.health = CaptureHealth.CONNECTED
            self._status.backend_name = backend_name
            self._status.consecutive_failures = 0
            self._reconnect_attempt = 0
        if self.config.verbose:
            st = self.status
            print(
                f"OpenCV negotiated stream: FOURCC='{st.fourcc}', "
                f"size={st.width}x{st.height}, fps={st.fps:.2f}, backend={st.backend_name}"
            )

    def release(self) -> None:
        with self._lock:
            self._release_unlocked()
            self._status.health = CaptureHealth.CLOSED

    def _release_unlocked(self) -> None:
        if self._cap is not None:
            try:
                self._cap.release()
            except Exception:
                pass
            self._cap = None

    def _refresh_stream_meta_unlocked(self) -> None:
        cap = self._cap
        if cap is None:
            return
        try:
            self._status.width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
            self._status.height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
            self._status.fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0)
            self._status.fourcc = _decode_fourcc(cap.get(cv2.CAP_PROP_FOURCC) or 0)
        except Exception:
            pass

    def read_raw(self) -> tuple[bool, Optional[np.ndarray]]:
        """Read one frame from the underlying capture (no reconnect)."""
        cap = self._cap
        if cap is None:
            return False, None
        try:
            ok, frame = cap.read()
        except Exception as exc:
            with self._lock:
                self._status.consecutive_failures += 1
                self._status.last_error = str(exc)
            return False, None
        if not ok or frame is None:
            with self._lock:
                self._status.consecutive_failures += 1
            return False, None
        with self._lock:
            self._status.consecutive_failures = 0
            self._status.last_frame_time = time.perf_counter()
            if self._status.health != CaptureHealth.CONNECTED:
                self._status.health = CaptureHealth.CONNECTED
        return True, frame

    def note_published(self, seq: int) -> None:
        with self._lock:
            self._status.last_seq = int(seq)

    def note_dropped(self, n: int = 1) -> None:
        with self._lock:
            self._status.dropped_frames += int(n)

    def should_reconnect(self) -> bool:
        policy = self.config.reconnect
        if not policy.enabled:
            return False
        with self._lock:
            if self._status.health == CaptureHealth.CLOSED:
                return False
            if not self.is_open():
                return True
            if self._status.consecutive_failures >= max(1, policy.fail_threshold):
                return True
            if policy.stale_timeout_s is not None and self._status.last_frame_time is not None:
                age = time.perf_counter() - self._status.last_frame_time
                if age >= policy.stale_timeout_s and self._status.health == CaptureHealth.CONNECTED:
                    return True
            return False

    def reconnect(self) -> bool:
        """Attempt to reopen the capture. Returns True on success."""
        policy: ReconnectPolicy = self.config.reconnect
        if not policy.enabled:
            return False
        if policy.max_attempts is not None and self._reconnect_attempt >= policy.max_attempts:
            with self._lock:
                self._status.health = CaptureHealth.FAILED
                self._status.last_error = "Reconnect attempts exhausted"
            return False

        with self._lock:
            self._status.health = CaptureHealth.RECONNECTING
            self._status.reconnect_attempts += 1
            attempt = self._reconnect_attempt
            self._reconnect_attempt += 1
            self._release_unlocked()

        delay = policy.next_backoff(attempt)
        if self.config.verbose:
            print(f"Reconnecting to {self.config.source!r} (attempt {attempt + 1}, wait {delay:.2f}s)...")
        time.sleep(delay)
        try:
            self.open()
            return True
        except Exception as exc:
            with self._lock:
                self._status.last_error = str(exc)
                if policy.max_attempts is not None and self._reconnect_attempt >= policy.max_attempts:
                    self._status.health = CaptureHealth.FAILED
                else:
                    self._status.health = CaptureHealth.RECONNECTING
            return False