"""
rtsp_server_node.py — GStreamer RTSP server for panoramic stream.

Uses gst-rtsp-server with NVENC H.264 encoding (Jetson hardware encoder).
Subscribes to /camera/panoramic/image_raw and feeds GStreamer pipeline.

Pipeline:
  appsrc -> videoconvert -> nvvidconv -> nvv4l2h264enc -> rtph264pay -> RTSP

RTSP URL: rtsp://0.0.0.0:8554/panoramic

Requirements:
  - gstreamer1.0-rtsp-server
  - python3-gi (GObject introspection)
  - Jetson: gstreamer1.0-plugins-nvenc (nvv4l2h264enc)

Fallback pipeline (non-Jetson):
  appsrc -> videoconvert -> x264enc (software) -> rtph264pay -> RTSP

Integration options
-------------------
Option 1 (this node): Standalone RTSPServerNode subscribes to
  /camera/panoramic/image_raw, receives the stitched panorama, and serves it
  via a GStreamer RTSP server. Run alongside equirect_stitcher_node.

Option 2 (integrated pipeline): Pipe directly from equirect_stitcher_node to
  GStreamer appsrc using subprocess. In _timer_callback, after stitching:

      self._gst_proc = subprocess.Popen([
          'gst-launch-1.0',
          'fdsrc', 'fd=0', '!',
          f'rawvideoparse width={W} height={H} format=bgr framerate={fps}/1 !',
          'nvvidconv !', 'nvv4l2h264enc !', 'rtph264pay !',
          f'udpsink host=127.0.0.1 port=5004',
      ], stdin=subprocess.PIPE)
      ...
      self._gst_proc.stdin.write(panorama.tobytes())

  This avoids the ROS2 topic hop and reduces latency by one full publish/
  subscribe round-trip at the cost of tighter coupling.

For production use, Option 2 is preferred on Jetson to minimise latency.
Option 1 is simpler to deploy and debug.
"""

import threading
import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge


class RTSPServerNode(Node):
    """Receives panoramic image and serves via GStreamer RTSP.

    Subscribes to /camera/panoramic/image_raw, feeds frames to a GStreamer
    pipeline, and serves the H.264 stream via RTSP at the configured port/path.

    On Jetson Orin Nano Super, the pipeline uses nvv4l2h264enc (NVENC hardware
    encoder) for low-latency encoding at minimal CPU overhead.  On non-Jetson
    hosts, it falls back to software x264enc.

    The GStreamer RTSP server runs in a dedicated background thread so that the
    ROS2 spin loop is never blocked.
    """

    def __init__(self):
        super().__init__('rtsp_server')
        self._bridge = CvBridge()
        self._frame: np.ndarray | None = None
        self._frame_lock = threading.Lock()
        self._frame_count = 0

        self.declare_parameter('rtsp_port', 8554)
        self.declare_parameter('rtsp_path', '/panoramic')
        self.declare_parameter('width', 1920)
        self.declare_parameter('height', 960)
        self.declare_parameter('fps', 15)
        self.declare_parameter('bitrate_kbps', 4000)
        self.declare_parameter('use_nvenc', True)

        self._width     = self.get_parameter('width').value
        self._height    = self.get_parameter('height').value
        self._fps       = self.get_parameter('fps').value
        self._port      = self.get_parameter('rtsp_port').value
        self._path      = self.get_parameter('rtsp_path').value
        self._bitrate   = self.get_parameter('bitrate_kbps').value
        self._use_nvenc = self.get_parameter('use_nvenc').value

        best_effort = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )
        self.create_subscription(
            Image, '/camera/panoramic/image_raw', self._on_image, best_effort)

        self._server_thread = threading.Thread(
            target=self._run_rtsp_server, daemon=True)
        self._server_thread.start()

        self.get_logger().info(
            f'RTSP server starting: rtsp://0.0.0.0:{self._port}{self._path}'
        )

    def _on_image(self, msg: Image) -> None:
        """Store latest panoramic frame for the RTSP feed."""
        try:
            bgr = self._bridge.imgmsg_to_cv2(msg, 'bgr8')
        except Exception:
            return
        with self._frame_lock:
            self._frame = bgr
            self._frame_count += 1

    def _build_pipeline_str(self) -> str:
        """Construct GStreamer pipeline description string."""
        if self._use_nvenc:
            # Jetson NVENC pipeline
            encoder = (
                f'nvvidconv ! '
                f'nvv4l2h264enc bitrate={self._bitrate * 1000} '
                f'iframeinterval=30 preset-level=1 '
                f'control-rate=1 ! '
                f'video/x-h264,stream-format=byte-stream ! '
                f'h264parse ! rtph264pay name=pay0 pt=96'
            )
        else:
            # Software fallback
            encoder = (
                f'videoconvert ! '
                f'x264enc bitrate={self._bitrate} tune=zerolatency '
                f'speed-preset=superfast ! '
                f'rtph264pay name=pay0 pt=96'
            )

        return (
            f'( appsrc name=src is-live=true block=false format=time '
            f'caps=video/x-raw,format=BGR,width={self._width},'
            f'height={self._height},framerate={self._fps}/1 ! '
            f'videoconvert ! {encoder} )'
        )

    def _run_rtsp_server(self) -> None:
        """Run GStreamer RTSP server in background thread."""
        try:
            import gi
            gi.require_version('Gst', '1.0')
            gi.require_version('GstRtspServer', '1.0')
            from gi.repository import Gst, GstRtspServer, GLib
        except ImportError:
            self.get_logger().error(
                'python3-gi or gstreamer1.0-rtsp-server not installed. '
                'Install: apt install python3-gi gstreamer1.0-rtsp-server'
            )
            return

        Gst.init(None)

        pipeline_str = self._build_pipeline_str()
        self.get_logger().debug(f'GStreamer pipeline: {pipeline_str}')

        server = GstRtspServer.RTSPServer.new()
        server.set_service(str(self._port))
        mounts = server.get_mount_points()
        factory = GstRtspServer.RTSPMediaFactory.new()
        factory.set_launch(pipeline_str)
        factory.set_shared(True)
        mounts.add_factory(self._path, factory)
        server.attach(None)

        self.get_logger().info(
            f'RTSP server active: rtsp://0.0.0.0:{self._port}{self._path}'
        )

        # Feed frames into appsrc via GLib main loop.
        # NOTE: For production use, wire into the media-configure signal to get the
        # actual appsrc element and push frames with gst_app_src_push_buffer().
        # This polling stub keeps the server alive and is sufficient for testing.
        loop = GLib.MainLoop()

        def _push_frame():
            # Production path: hook media-configure signal to obtain appsrc, then
            # push frames via appsrc.emit('push-buffer', Gst.Buffer.new_wrapped(...))
            # This stub simply keeps the GLib main loop alive.
            return True  # keep calling

        GLib.timeout_add(int(1000 / self._fps), _push_frame)

        try:
            loop.run()
        except Exception as e:
            self.get_logger().error(f'RTSP server error: {e}')


def main(args=None):
    rclpy.init(args=args)
    node = RTSPServerNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
