#!/usr/bin/env python3
"""
phone_camera_node.py — Jetson ROS2 receiver for phone video stream (Issue #585)

Connects to the phone's video_bridge.py WebSocket server, receives JPEG frames,
decodes them, and publishes:
  /saltybot/phone/camera            sensor_msgs/Image            (bgr8)
  /saltybot/phone/camera/compressed sensor_msgs/CompressedImage  (jpeg)

Falls back to HTTP MJPEG polling if the WebSocket connection fails.

Parameters
──────────
  phone_host      str   Phone IP or hostname     (default: 192.168.1.200)
  ws_port         int   video_bridge WS port      (default: 8765)
  http_port       int   video_bridge HTTP port    (default: 8766)
  frame_id        str   TF frame id               (default: phone_camera)
  publish_raw     bool  Publish sensor_msgs/Image (default: true)
  publish_compressed bool Publish CompressedImage (default: true)
  reconnect_s     float Seconds between reconnect (default: 3.0)
  latency_warn_ms float Warn if frame age > this  (default: 200.0)

Publishes
─────────
  /saltybot/phone/camera            sensor_msgs/Image
  /saltybot/phone/camera/compressed sensor_msgs/CompressedImage
  /saltybot/phone/camera/info       std_msgs/String  (JSON stream metadata)
"""

import asyncio
import json
import threading
import time
from typing import Optional

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image, CompressedImage
from std_msgs.msg import String, Header

try:
    import cv2
    import numpy as np
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False

try:
    import websockets
    WS_AVAILABLE = True
except ImportError:
    WS_AVAILABLE = False

try:
    import urllib.request
    HTTP_AVAILABLE = True
except ImportError:
    HTTP_AVAILABLE = False


MJPEG_BOUNDARY = b"--mjpeg-boundary"


class PhoneCameraNode(Node):
    """
    Receives JPEG frames from phone/video_bridge.py and publishes ROS2 Image topics.

    Connection strategy:
      1. Try WebSocket (ws://<phone_host>:<ws_port>)  — lowest latency
      2. On WS failure, fall back to HTTP MJPEG stream — guaranteed to work
      3. Reconnect automatically on disconnect
    """

    def __init__(self):
        super().__init__("phone_camera_node")

        # Parameters
        self.declare_parameter("phone_host",          "192.168.1.200")
        self.declare_parameter("ws_port",             8765)
        self.declare_parameter("http_port",           8766)
        self.declare_parameter("frame_id",            "phone_camera")
        self.declare_parameter("publish_raw",         True)
        self.declare_parameter("publish_compressed",  True)
        self.declare_parameter("reconnect_s",         3.0)
        self.declare_parameter("latency_warn_ms",     200.0)

        self._host         = self.get_parameter("phone_host").value
        self._ws_port      = self.get_parameter("ws_port").value
        self._http_port    = self.get_parameter("http_port").value
        self._frame_id     = self.get_parameter("frame_id").value
        self._pub_raw      = self.get_parameter("publish_raw").value
        self._pub_comp     = self.get_parameter("publish_compressed").value
        self._reconnect_s  = self.get_parameter("reconnect_s").value
        self._latency_warn = self.get_parameter("latency_warn_ms").value

        if not CV2_AVAILABLE:
            self.get_logger().error(
                "opencv-python not installed — cannot decode JPEG frames."
            )

        # Publishers
        self._raw_pub  = self.create_publisher(Image,           "/saltybot/phone/camera",            10)
        self._comp_pub = self.create_publisher(CompressedImage, "/saltybot/phone/camera/compressed",  10)
        self._info_pub = self.create_publisher(String,          "/saltybot/phone/camera/info",        10)

        # Stats
        self._frames_received  = 0
        self._frames_published = 0
        self._last_frame_ts    = 0.0
        self._stream_info: dict = {}

        # Diagnostics timer (every 10 s)
        self.create_timer(10.0, self._diag_cb)

        # Start receiver thread
        self._stop_event = threading.Event()
        self._recv_thread = threading.Thread(
            target=self._recv_loop, name="phone_cam_recv", daemon=True
        )
        self._recv_thread.start()

        self.get_logger().info(
            "PhoneCameraNode ready — connecting to %s (ws:%d / http:%d)",
            self._host, self._ws_port, self._http_port,
        )

    # ── Publish ───────────────────────────────────────────────────────────────

    def _publish_jpeg(self, jpeg_bytes: bytes, recv_ts: float) -> None:
        """Decode JPEG and publish Image and/or CompressedImage."""
        self._frames_received += 1
        age_ms = (time.time() - recv_ts) * 1000.0
        if age_ms > self._latency_warn:
            self.get_logger().warning(
                "Phone camera frame age %.0f ms (target < %.0f ms)",
                age_ms, self._latency_warn,
            )

        now_msg = self.get_clock().now().to_msg()

        # CompressedImage — no decode needed
        if self._pub_comp:
            comp = CompressedImage()
            comp.header = Header()
            comp.header.stamp    = now_msg
            comp.header.frame_id = self._frame_id
            comp.format = "jpeg"
            comp.data   = list(jpeg_bytes)
            self._comp_pub.publish(comp)

        # Image — decode via OpenCV
        if self._pub_raw and CV2_AVAILABLE:
            try:
                buf   = np.frombuffer(jpeg_bytes, dtype=np.uint8)
                frame = cv2.imdecode(buf, cv2.IMREAD_COLOR)
                if frame is None:
                    return
                h, w = frame.shape[:2]
                img = Image()
                img.header        = Header()
                img.header.stamp    = now_msg
                img.header.frame_id = self._frame_id
                img.height          = h
                img.width           = w
                img.encoding        = "bgr8"
                img.is_bigendian    = 0
                img.step            = w * 3
                img.data            = frame.tobytes()
                self._raw_pub.publish(img)
            except Exception as e:
                self.get_logger().debug("JPEG decode error: %s", str(e))
                return

        self._frames_published += 1
        self._last_frame_ts = time.time()

    # ── WebSocket receiver ────────────────────────────────────────────────────

    async def _ws_recv(self) -> bool:
        """
        Connect via WebSocket and receive frames until disconnect.
        Returns True if we received at least one frame (connection was working).
        """
        uri = f"ws://{self._host}:{self._ws_port}"
        try:
            async with websockets.connect(
                uri,
                ping_interval=5,
                ping_timeout=10,
                max_size=None,
                open_timeout=5,
            ) as ws:
                self.get_logger().info("WebSocket connected to %s", uri)
                got_frame = False
                async for message in ws:
                    if self._stop_event.is_set():
                        return got_frame
                    recv_ts = time.time()
                    if isinstance(message, bytes):
                        # Binary → JPEG frame
                        self._publish_jpeg(message, recv_ts)
                        got_frame = True
                    elif isinstance(message, str):
                        # Text → control JSON
                        try:
                            obj = json.loads(message)
                            if obj.get("type") == "info":
                                self._stream_info = obj
                                info_msg = String()
                                info_msg.data = message
                                self._info_pub.publish(info_msg)
                                self.get_logger().info(
                                    "Stream info: %dx%d @ %.0f fps",
                                    obj.get("width", 0),
                                    obj.get("height", 0),
                                    obj.get("fps", 0),
                                )
                        except json.JSONDecodeError:
                            pass
                return got_frame
        except (websockets.exceptions.WebSocketException, OSError,
                asyncio.TimeoutError) as e:
            self.get_logger().debug("WS error: %s", str(e))
            return False

    # ── HTTP MJPEG fallback ───────────────────────────────────────────────────

    def _http_mjpeg_recv(self) -> None:
        """Receive MJPEG stream over HTTP and publish frames."""
        url = f"http://{self._host}:{self._http_port}/stream"
        self.get_logger().info("Falling back to HTTP MJPEG: %s", url)
        buf = b""
        try:
            req = urllib.request.urlopen(url, timeout=10)
            while not self._stop_event.is_set():
                chunk = req.read(4096)
                if not chunk:
                    break
                buf += chunk
                # Extract JPEG frames delimited by MJPEG boundary
                while True:
                    start = buf.find(b"\xff\xd8")  # JPEG SOI
                    end   = buf.find(b"\xff\xd9")  # JPEG EOI
                    if start == -1 or end == -1 or end < start:
                        # Keep last partial frame in buffer
                        if start > 0:
                            buf = buf[start:]
                        break
                    jpeg = buf[start:end + 2]
                    buf  = buf[end + 2:]
                    self._publish_jpeg(jpeg, time.time())
        except Exception as e:
            self.get_logger().debug("HTTP MJPEG error: %s", str(e))

    # ── Receiver loop ─────────────────────────────────────────────────────────

    def _recv_loop(self) -> None:
        """
        Main receiver loop — tries WS first, falls back to HTTP, then retries.
        Runs in a daemon thread.
        """
        while not self._stop_event.is_set():
            # Try WebSocket
            if WS_AVAILABLE:
                try:
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                    got = loop.run_until_complete(self._ws_recv())
                    loop.close()
                    if got and not self._stop_event.is_set():
                        # WS was working but disconnected — retry immediately
                        continue
                except Exception as e:
                    self.get_logger().debug("WS loop error: %s", str(e))

            if self._stop_event.is_set():
                break

            # WS failed — try HTTP MJPEG fallback
            if HTTP_AVAILABLE:
                self._http_mjpeg_recv()

            if not self._stop_event.is_set():
                self.get_logger().info(
                    "Phone camera stream lost — retrying in %.0f s", self._reconnect_s
                )
                self._stop_event.wait(self._reconnect_s)

    # ── Diagnostics ───────────────────────────────────────────────────────────

    def _diag_cb(self) -> None:
        age = time.time() - self._last_frame_ts if self._last_frame_ts > 0 else -1.0
        self.get_logger().info(
            "Phone camera — received=%d published=%d last_frame_age=%.1fs",
            self._frames_received, self._frames_published, age,
        )

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    def destroy_node(self) -> None:
        self._stop_event.set()
        super().destroy_node()


def main(args=None) -> None:
    rclpy.init(args=args)
    node = PhoneCameraNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()
