"""Frame capture utility for the Edge Proxy.

Grabs a single JPEG frame from the robot's camera.  The primary method is an
HTTP snapshot endpoint exposed by MediaMTX.  If that endpoint is unavailable,
a **persistent** ffmpeg subprocess maintains an RTSP connection and continuously
decodes frames into JPEG.  Each ``capture()`` call returns the most recently
decoded frame without re-negotiating the RTSP session.
"""

from __future__ import annotations

import io
import logging
import os
import subprocess
import threading
import time
import urllib.error
import urllib.request
from typing import Optional

logger = logging.getLogger(__name__)

# Default URLs for the Ghost PC's MediaMTX instance on the robot subnet.
_DEFAULT_SNAPSHOT_URL = "http://192.168.168.105:8889/cam/jpeg"
_DEFAULT_RTSP_URL = "rtsp://192.168.168.105:8554/cam"

_HTTP_TIMEOUT_SEC = 5

# JPEG markers
_SOI = b"\xff\xd8"
_EOI = b"\xff\xd9"


class FrameCapture:
    """Capture a single JPEG frame from the robot camera.

    Attempts an HTTP snapshot first (fast, ~100 ms).  Falls back to a
    **persistent** ffmpeg subprocess that holds the RTSP connection open and
    continuously pipes JPEG frames via stdout.  The latest frame is always
    available in ``capture()`` with near-zero latency.

    Args:
        snapshot_url: HTTP URL that returns a raw JPEG.  MediaMTX exposes this
            at ``/cam/jpeg`` when the API is enabled.
        rtsp_url: RTSP URL used for the persistent ffmpeg stream.
    """

    def __init__(
        self,
        snapshot_url: str = _DEFAULT_SNAPSHOT_URL,
        rtsp_url: str = _DEFAULT_RTSP_URL,
    ) -> None:
        self._snapshot_url = snapshot_url
        self._rtsp_url = rtsp_url

        # Persistent RTSP stream state
        self._proc: Optional[subprocess.Popen] = None
        self._reader_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()
        self._latest_frame: Optional[bytes] = None
        self._frame_time: float = 0.0
        self._frame_wall_time: float = 0.0
        self._running = False
        self._http_ok: Optional[bool] = None  # None = untested
        self._last_capture_ms: float = 0.0
        self._last_capture_source: str = "none"

        # Probe HTTP once at init (on main thread, before asyncio starts)
        self._http_ok = self._capture_http() is not None
        if self._http_ok:
            logger.info("FrameCapture: HTTP snapshot available — using fast path")
        else:
            logger.info(
                "FrameCapture: HTTP snapshot unavailable — starting persistent RTSP stream"
            )
            self._ensure_rtsp_stream()

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

    def capture(self) -> Optional[bytes]:
        """Grab a single JPEG frame.

        Returns:
            Raw JPEG bytes, or ``None`` if all methods fail.
        """
        t0 = time.monotonic()

        # Fast path: HTTP snapshot
        if self._http_ok:
            frame = self._capture_http()
            if frame is not None:
                self._update_capture_stats(
                    source="http",
                    capture_ms=(time.monotonic() - t0) * 1000.0,
                )
                logger.debug("FrameCapture: HTTP snapshot succeeded (%d bytes)", len(frame))
                return frame

        # Persistent RTSP path — stream already started in __init__
        # Re-start if ffmpeg died
        self._ensure_rtsp_stream()
        with self._lock:
            frame = self._latest_frame
            frame_age_ms = (
                (time.monotonic() - self._frame_time) * 1000.0
                if self._frame_time > 0
                else None
            )
        self._update_capture_stats(
            source="rtsp",
            capture_ms=(time.monotonic() - t0) * 1000.0,
        )
        if frame is not None:
            if frame_age_ms is not None:
                logger.debug(
                    "FrameCapture: RTSP frame ready (%d bytes, age=%.1f ms)",
                    len(frame),
                    frame_age_ms,
                )
            else:
                logger.debug("FrameCapture: RTSP frame ready (%d bytes)", len(frame))
        else:
            logger.debug("FrameCapture: RTSP stream not yet producing frames")
        return frame

    def get_stats(self) -> dict:
        """Return current frame/capture stats for latency debugging."""
        with self._lock:
            frame_age_ms = (
                (time.monotonic() - self._frame_time) * 1000.0
                if self._frame_time > 0
                else None
            )
            latest_frame_ts = self._frame_wall_time
            last_capture_ms = self._last_capture_ms
            last_capture_source = self._last_capture_source
        return {
            "frame_age_ms": frame_age_ms,
            "latest_frame_timestamp": latest_frame_ts,
            "last_capture_ms": last_capture_ms,
            "last_capture_source": last_capture_source,
            "rtsp_running": bool(
                self._running and self._proc is not None and self._proc.poll() is None
            ),
            "http_snapshot_ok": self._http_ok,
        }

    def close(self) -> None:
        """Shut down the persistent ffmpeg subprocess."""
        self._running = False
        if self._proc is not None:
            try:
                self._proc.terminate()
                self._proc.wait(timeout=3)
            except Exception:  # noqa: BLE001
                self._proc.kill()
            self._proc = None
        if self._reader_thread is not None:
            self._reader_thread.join(timeout=3)
            self._reader_thread = None
        logger.info("FrameCapture: persistent RTSP stream closed")

    # ------------------------------------------------------------------
    # HTTP snapshot (fast path)
    # ------------------------------------------------------------------

    def _capture_http(self) -> Optional[bytes]:
        """Try HTTP snapshot endpoint.  Returns raw bytes or None."""
        try:
            with urllib.request.urlopen(self._snapshot_url, timeout=_HTTP_TIMEOUT_SEC) as resp:
                if resp.status != 200:
                    logger.debug(
                        "FrameCapture: HTTP snapshot returned status %d", resp.status
                    )
                    return None
                data = resp.read()
                if not data:
                    logger.debug("FrameCapture: HTTP snapshot returned empty body")
                    return None
                return data
        except urllib.error.URLError as exc:
            logger.debug("FrameCapture: HTTP snapshot URL error: %s", exc)
            return None
        except OSError as exc:
            logger.debug("FrameCapture: HTTP snapshot OS error: %s", exc)
            return None
        except Exception as exc:  # noqa: BLE001
            logger.debug("FrameCapture: HTTP snapshot unexpected error: %s", exc)
            return None

    # ------------------------------------------------------------------
    # Persistent RTSP stream (ffmpeg subprocess)
    # ------------------------------------------------------------------

    def _ensure_rtsp_stream(self) -> None:
        """Start the persistent ffmpeg process if not already running."""
        if self._running and self._proc is not None and self._proc.poll() is None:
            return  # already running

        # Clean up any dead process
        if self._proc is not None:
            try:
                self._proc.kill()
            except Exception:  # noqa: BLE001
                pass
            self._proc = None

        logger.info("FrameCapture: starting persistent RTSP stream from %s", self._rtsp_url)
        try:
            self._proc = subprocess.Popen(
                [
                    "ffmpeg",
                    "-fflags", "nobuffer",
                    "-flags", "low_delay",
                    "-avioflags", "direct",
                    "-probesize", "32",
                    "-analyzeduration", "0",
                    "-rtsp_transport", "tcp",
                    "-i", self._rtsp_url,
                    "-vf", "fps=10",                 # match camera native 10fps
                    "-f", "image2pipe",
                    "-vcodec", "mjpeg",
                    "-q:v", "3",
                    "-",
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.DEVNULL,
                bufsize=0,
            )
        except FileNotFoundError:
            logger.warning("FrameCapture: ffmpeg not found in PATH")
            return

        self._running = True
        self._reader_thread = threading.Thread(
            target=self._read_frames, daemon=True, name="frame-capture-reader"
        )
        self._reader_thread.start()

    def _read_frames(self) -> None:
        """Background thread: read continuous JPEG stream from ffmpeg stdout.

        ffmpeg with ``-f image2pipe -vcodec mjpeg`` outputs concatenated JPEG
        images.  We split on SOI/EOI markers to extract individual frames.
        """
        try:
            assert self._proc is not None and self._proc.stdout is not None
            buf = bytearray()
            stream = self._proc.stdout
            frame_count = 0

            while self._running:
                try:
                    chunk = stream.read(65536)
                except Exception as exc:  # noqa: BLE001
                    logger.warning("FrameCapture: read error: %s", exc)
                    break
                if not chunk:
                    break  # EOF — ffmpeg exited

                buf.extend(chunk)

                # Extract the latest complete JPEG frame from the buffer
                while True:
                    soi = buf.find(_SOI)
                    if soi == -1:
                        buf.clear()
                        break

                    # Discard anything before SOI
                    if soi > 0:
                        del buf[:soi]

                    eoi = buf.find(_EOI, 2)  # search after SOI
                    if eoi == -1:
                        break  # incomplete frame, wait for more data

                    # Extract complete JPEG (SOI through EOI inclusive)
                    frame_end = eoi + 2
                    frame = bytes(buf[:frame_end])
                    del buf[:frame_end]
                    frame_count += 1

                    with self._lock:
                        self._latest_frame = frame
                        self._frame_time = time.monotonic()
                        self._frame_wall_time = time.time()

                    if frame_count == 1:
                        logger.info(
                            "FrameCapture: first RTSP frame received (%d bytes)", len(frame)
                        )

        except Exception as exc:
            logger.error("FrameCapture: reader thread crashed: %s", exc, exc_info=True)
        finally:
            self._running = False
            logger.info("FrameCapture: RTSP reader thread exited (frames read: %d)",
                        frame_count if 'frame_count' in dir() else 0)

    def _update_capture_stats(self, source: str, capture_ms: float) -> None:
        with self._lock:
            self._last_capture_source = source
            self._last_capture_ms = capture_ms