"""
rtsp_stream.stream
~~~~~~~~~~~~~~~~~~
Thread-based RTSP video stream reader with automatic reconnection.
"""

from __future__ import annotations

import logging
import threading
import time
from datetime import datetime
from typing import Callable, Optional

import cv2
import numpy as np

logger = logging.getLogger(__name__)


class VideoStream:
    """
    Thread-based RTSP (or any OpenCV-compatible) video stream reader
    with automatic reconnection on feed drops.

    Parameters
    ----------
    rtsp_url : str
        RTSP URL (or any source accepted by ``cv2.VideoCapture``).
    reconnect_delay : float
        Seconds to wait between reconnection attempts. Default: ``1.0``.
    max_reconnect_attempts : int | None
        Maximum consecutive reconnection attempts before giving up.
        ``None`` (default) means retry forever.
    reconnect_backoff : float
        Multiplicative back-off factor applied to *reconnect_delay* after
        each failed attempt (capped at *max_reconnect_delay*). Default: ``1.0``
        (no back-off).
    max_reconnect_delay : float
        Upper bound on the reconnect sleep interval when using back-off.
        Default: ``30.0`` seconds.
    on_connect : Callable | None
        Optional callback fired every time a connection (or re-connection)
        succeeds.  Receives no arguments.
    on_disconnect : Callable | None
        Optional callback fired every time the feed is lost.  Receives no
        arguments.
    on_reconnect_failure : Callable | None
        Optional callback fired when ``max_reconnect_attempts`` is exhausted.
        Receives no arguments.
    cap_backend : int
        OpenCV backend flag passed to ``cv2.VideoCapture``.
        Default: ``cv2.CAP_ANY``.
    cap_buffer_size : int | None
        If set, configures ``cv2.CAP_PROP_BUFFERSIZE`` right after opening.
        Useful for reducing latency on RTSP streams.
    """

    def __init__(
        self,
        rtsp_url: str,
        *,
        reconnect_delay: float = 1.0,
        max_reconnect_attempts: Optional[int] = None,
        reconnect_backoff: float = 1.0,
        max_reconnect_delay: float = 30.0,
        on_connect: Optional[Callable[[], None]] = None,
        on_disconnect: Optional[Callable[[], None]] = None,
        on_reconnect_failure: Optional[Callable[[], None]] = None,
        cap_backend: int = cv2.CAP_ANY,
        cap_buffer_size: Optional[int] = None,
    ) -> None:
        self.rtsp_url = rtsp_url
        self.reconnect_delay = reconnect_delay
        self.max_reconnect_attempts = max_reconnect_attempts
        self.reconnect_backoff = reconnect_backoff
        self.max_reconnect_delay = max_reconnect_delay
        self.cap_backend = cap_backend
        self.cap_buffer_size = cap_buffer_size

        # Callbacks
        self._on_connect = on_connect
        self._on_disconnect = on_disconnect
        self._on_reconnect_failure = on_reconnect_failure

        # Internal state
        self._frame: Optional[np.ndarray] = None
        self._frame_lock = threading.Lock()
        self._running = False
        self._connected = False
        self._cap: Optional[cv2.VideoCapture] = None
        self._reconnect_count = 0

        self._thread = threading.Thread(target=self._run, daemon=True)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def start(self) -> "VideoStream":
        """Start the background capture thread.  Returns *self* for chaining."""
        if self._running:
            raise RuntimeError("VideoStream is already running.")
        self._running = True
        self._thread.start()
        logger.debug("VideoStream started for %s", self.rtsp_url)
        return self

    def read(self) -> Optional[np.ndarray]:
        """Return the most recent frame, or ``None`` if none is available yet."""
        with self._frame_lock:
            return self._frame

    def stop(self) -> None:
        """Signal the capture thread to stop and block until it exits."""
        self._running = False
        self._thread.join(timeout=10)
        self._release_cap()
        logger.debug("VideoStream stopped.")

    @property
    def is_connected(self) -> bool:
        """``True`` while a capture object is open and delivering frames."""
        return self._connected

    @property
    def reconnect_count(self) -> int:
        """Total number of reconnection attempts made since :meth:`start`."""
        return self._reconnect_count

    # Capture properties (safe to call at any time; return 0.0 when offline)

    @property
    def fps(self) -> float:
        return self._cap_get(cv2.CAP_PROP_FPS)

    @property
    def frame_width(self) -> int:
        return int(self._cap_get(cv2.CAP_PROP_FRAME_WIDTH))

    @property
    def frame_height(self) -> int:
        return int(self._cap_get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Context-manager support — use as: `with VideoStream(...).start() as stream:`
    def __enter__(self) -> "VideoStream":
        if not self._running:
            self.start()
        return self

    def __exit__(self, *_) -> None:
        self.stop()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _cap_get(self, prop: int) -> float:
        return self._cap.get(prop) if self._cap and self._cap.isOpened() else 0.0

    def _release_cap(self) -> None:
        if self._cap:
            self._cap.release()
            self._cap = None

    def _open_cap(self) -> bool:
        """Try to open the capture.  Returns True on success."""
        cap = cv2.VideoCapture(self.rtsp_url, self.cap_backend)
        if not cap.isOpened():
            cap.release()
            return False
        if self.cap_buffer_size is not None:
            cap.set(cv2.CAP_PROP_BUFFERSIZE, self.cap_buffer_size)
        self._cap = cap
        return True

    def _connect_with_retry(self) -> bool:
        """
        Block until connected or until max_reconnect_attempts is reached.
        Returns True if a connection was established, False otherwise.
        """
        delay = self.reconnect_delay
        attempt = 0

        while self._running:
            if self._open_cap():
                self._reconnect_count += attempt  # only count failed tries
                self._connected = True
                logger.info("Connected to %s", self.rtsp_url)
                if self._on_connect:
                    try:
                        self._on_connect()
                    except Exception:  # noqa: BLE001
                        logger.exception("on_connect callback raised an exception")
                return True

            attempt += 1
            logger.warning(
                "Connection attempt %d failed for %s – retrying in %.1fs",
                attempt,
                self.rtsp_url,
                delay,
            )

            if self.max_reconnect_attempts is not None and attempt >= self.max_reconnect_attempts:
                logger.error(
                    "Exhausted %d reconnect attempts for %s",
                    self.max_reconnect_attempts,
                    self.rtsp_url,
                )
                if self._on_reconnect_failure:
                    try:
                        self._on_reconnect_failure()
                    except Exception:  # noqa: BLE001
                        logger.exception("on_reconnect_failure callback raised an exception")
                return False

            time.sleep(delay)
            delay = min(delay * self.reconnect_backoff, self.max_reconnect_delay)

        return False  # stopped while waiting

    def _run(self) -> None:
        """Main loop running in the daemon thread."""
        if not self._connect_with_retry():
            self._running = False
            return

        while self._running:
            ret, frame = self._cap.read()

            if ret:
                with self._frame_lock:
                    self._frame = frame
            else:
                # Feed dropped
                self._connected = False
                ts = int(datetime.now().timestamp() * 1000)
                logger.warning("Feed dropped for %s at %d ms", self.rtsp_url, ts)

                with self._frame_lock:
                    self._frame = None

                self._release_cap()

                if self._on_disconnect:
                    try:
                        self._on_disconnect()
                    except Exception:  # noqa: BLE001
                        logger.exception("on_disconnect callback raised an exception")

                if not self._connect_with_retry():
                    self._running = False
                    return

        self._connected = False
