"""Hosts rclpy inside the FastAPI process.

Mirrors the isolation pattern used by RosSubscriber in scan_gui_triple.py:
a dedicated rclpy.Context + Node, spun in a daemon thread, kept out of the
ASGI event loop entirely. One node, one MultiThreadedExecutor, one shared
ReentrantCallbackGroup for every subscription.

/livox/imu (~200Hz) is deliberately NOT subscribed here — it's ~20x every
other topic's rate, and every way of grouping it alongside slower topics
(shared default group, shared Reentrant group, its own dedicated
Context+Node+thread, split-by-cost with the other cheap raw subscriptions)
measurably starved the slower ones, cameras worst of all — confirmed
2026-08-06 across several configurations, right down to cameras measuring
~0Hz against a real ~10Hz feed even after every other topic's callback was
made cheap (raw=True). The operator decided camera accuracy matters more
than a live IMU Hz reading, so IMU health is now inferred from lidar's
health instead of measured directly (see routers/system.py) — they come off
the same livox_ros_driver2 process, so lidar being alive is already strong
evidence IMU is too. If IMU health/Hz needs to be real again later, give it
back its own dedicated Context+Node+executor+thread (not a shared one) and
expect to spend real GIL budget on it.

Subscriptions cover: camera images (JPEG-encoded for MJPEG serving), the raw
lidar driver topic (health only), and (M3) FAST-LIO's point cloud +
odometry for the live scan view. GPS (M2) is not yet added. All
subscriptions register once at startup, never per start/stop call, so
health.py can tell "process down" apart from "never subscribed".

Cross-thread handoff: ROS callbacks (background thread) write into a
lock-guarded "latest value" dict; HTTP/WS handlers (asyncio loop) read it.
No queues — a slow reader just sees a slightly stale value, never backs up
ROS callback processing.
"""
from __future__ import annotations

import logging
import threading
import time
from typing import Optional

import config
from config import camera_specs
from health import TopicHealthTracker
from pointcloud_codec import PointCloudAccumulator

logger = logging.getLogger("scan_web.ros_bridge")

try:
    import numpy as np
    import rclpy
    from rclpy.callback_groups import ReentrantCallbackGroup
    from rclpy.executors import MultiThreadedExecutor
    from rclpy.qos import QoSProfile, QoSReliabilityPolicy, QoSHistoryPolicy
    from sensor_msgs.msg import Image, PointCloud2
    from nav_msgs.msg import Odometry
    from livox_ros_driver2.msg import CustomMsg
    from cv_bridge import CvBridge
    import sensor_msgs_py.point_cloud2 as pc2
    import cv2
    ROS_AVAILABLE = True
except ImportError:
    ROS_AVAILABLE = False


class ROSBridge:
    def __init__(self, camera_count: int):
        self.camera_count = camera_count
        self.health = TopicHealthTracker()
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._ctx = None  # set in _spin(); stop() shuts it down to unblock executor.spin()
        self._bridge = CvBridge() if ROS_AVAILABLE else None
        self._frame_lock = threading.Lock()
        self._latest_jpeg: dict[str, bytes] = {}
        self._raw_lock = threading.Lock()
        self._latest_raw: dict[str, object] = {}
        self._frame_events: dict[str, threading.Event] = {}
        self.pointcloud = PointCloudAccumulator(
            config.POINTCLOUD_VOXEL_SIZE_M, config.POINTCLOUD_MAX_ACCUM_POINTS,
        )

    # ── lifecycle ──────────────────────────────────────────────────────────

    def start(self) -> None:
        if not ROS_AVAILABLE:
            logger.warning("rclpy not available (ROS overlays not sourced?) — camera preview disabled")
            return
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(target=self._spin, daemon=True, name="ros-bridge")
        self._thread.start()

    def stop(self) -> None:
        self._running = False
        if self._ctx is not None:
            # Unblocks the executor.spin() call in _spin()'s thread from out
            # here — spin() only checks context.ok()/is_shutdown, it doesn't
            # know about self._running.
            try:
                self._ctx.shutdown()
            except Exception:
                pass
        if self._thread:
            self._thread.join(timeout=2.0)
            self._thread = None

    def _spin(self) -> None:
        ctx = rclpy.Context()
        ctx.init()
        self._ctx = ctx
        try:
            node = rclpy.create_node("scan_web_bridge", context=ctx)
            # One shared ReentrantCallbackGroup for every subscription below
            # — the node's default MutuallyExclusiveCallbackGroup would
            # serialize dispatch across all of them regardless of executor
            # type. No /livox/imu here (see module docstring) — with it
            # removed, the remaining topics are all roughly the same order
            # of magnitude (~10Hz), so they coexist fairly under one
            # Reentrant group without needing to be split further.
            cbg = ReentrantCallbackGroup()
            img_qos = QoSProfile(
                reliability=QoSReliabilityPolicy.BEST_EFFORT,
                history=QoSHistoryPolicy.KEEP_LAST,
                depth=1,
            )
            for cam in camera_specs(self.camera_count):
                node.create_subscription(
                    Image, cam.image_topic,
                    lambda msg, c=cam.id, t=cam.image_topic: self._on_image(c, t, msg),
                    img_qos, callback_group=cbg,
                )
            # Raw (undeserialized) lidar driver topic — subscribed only for
            # staleness/health tracking, not decoded/stored. This callback
            # never touches the payload, so raw=True skips deserialization
            # cost entirely (mirrors what `ros2 bag record` does
            # internally). Distinct from the LIO mapping output
            # (/aft_mapped_to_init etc.), which is M3 scope.
            sensor_qos = QoSProfile(
                reliability=QoSReliabilityPolicy.BEST_EFFORT,
                history=QoSHistoryPolicy.KEEP_LAST,
                depth=1,
            )
            node.create_subscription(
                CustomMsg, "/livox/lidar",
                lambda msg: self.health.mark_received("/livox/lidar"), sensor_qos, callback_group=cbg, raw=True,
            )
            # M3 — FAST-LIO's live mapping output. /cloud_registered is only
            # published while the "fastlio" process is running (tied to the
            # recording lifecycle, see routers/recording.py), same as lidar
            # is gated on the "lidar" process.
            node.create_subscription(
                PointCloud2, config.POINTCLOUD_TOPIC, self._on_pointcloud, sensor_qos, callback_group=cbg,
            )
            node.create_subscription(
                Odometry, config.ODOMETRY_TOPIC, self._on_odometry, sensor_qos, callback_group=cbg,
            )
            # MultiThreadedExecutor (not Single) — spin_once() on either type
            # still only pulls ONE ready entity per call, but Single runs its
            # callback synchronously inline before the next pull, whereas
            # MultiThreaded submits it to a worker pool and returns
            # immediately, so the poll loop keeps up with several ~10Hz
            # topics instead of falling behind and starving each other.
            # num_threads is deliberately just above the subscription count
            # (6: 3 cameras + lidar + pointcloud + odometry) rather than
            # generously oversized — Python's GIL means more threads doesn't
            # mean more parallel throughput, only more context-switch
            # overhead once you're past "enough to avoid queueing".
            executor = MultiThreadedExecutor(context=ctx, num_threads=6)
            executor.add_node(node)
            executor.spin()
            executor.remove_node(node)
            node.destroy_node()
        except Exception:
            logger.exception("ros bridge spin loop crashed")
        finally:
            self._ctx = None
            try:
                ctx.shutdown()
            except Exception:
                pass

    # ── callbacks ──────────────────────────────────────────────────────────

    def _on_image(self, cam_id: str, topic: str, msg) -> None:
        # Record arrival immediately and hand the deserialized msg off to a
        # per-camera worker thread. cv_bridge/JPEG encoding is too slow to
        # run inline here — running it directly in the ROS callback would
        # tie up that callback's executor thread for the encode's duration,
        # needlessly limiting how many camera frames can be in flight
        # concurrently.
        self.health.mark_received(topic)
        with self._raw_lock:
            self._latest_raw[cam_id] = msg
        ev = self._frame_events.get(cam_id)
        if ev is None:
            ev = threading.Event()
            self._frame_events[cam_id] = ev
            threading.Thread(
                target=self._encode_loop, args=(cam_id,), daemon=True, name=f"encode-{cam_id}",
            ).start()
        ev.set()

    def _encode_loop(self, cam_id: str) -> None:
        ev = self._frame_events[cam_id]
        while self._running:
            if not ev.wait(timeout=0.5):
                continue
            ev.clear()
            with self._raw_lock:
                msg = self._latest_raw.get(cam_id)
            if msg is None:
                continue
            try:
                frame = self._bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
                ok, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
                if not ok:
                    continue
                with self._frame_lock:
                    self._latest_jpeg[cam_id] = buf.tobytes()
            except Exception:
                logger.exception("failed to decode frame for %s", cam_id)

    def _on_pointcloud(self, msg) -> None:
        # Same reasoning as _on_image: read_points_numpy + voxel-dedup
        # accumulation is too slow to run inline on the shared spin thread,
        # so hand the raw msg off to a dedicated worker thread.
        self.health.mark_received(config.POINTCLOUD_TOPIC)
        with self._raw_lock:
            self._latest_raw["pointcloud"] = msg
        ev = self._frame_events.get("pointcloud")
        if ev is None:
            ev = threading.Event()
            self._frame_events["pointcloud"] = ev
            threading.Thread(target=self._pointcloud_loop, daemon=True, name="pointcloud-accum").start()
        ev.set()

    def _pointcloud_loop(self) -> None:
        ev = self._frame_events["pointcloud"]
        while self._running:
            if not ev.wait(timeout=0.5):
                continue
            ev.clear()
            with self._raw_lock:
                msg = self._latest_raw.get("pointcloud")
            if msg is None:
                continue
            try:
                arr = pc2.read_points_numpy(msg, field_names=("x", "y", "z", "intensity"), skip_nans=True)
                if arr.size == 0:
                    continue
                xyz = np.ascontiguousarray(arr[:, :3], dtype=np.float32)
                intensity = np.ascontiguousarray(arr[:, 3], dtype=np.float32)
                self.pointcloud.add_scan(xyz, intensity)
            except Exception:
                logger.exception("failed to accumulate point cloud")

    def _on_odometry(self, msg) -> None:
        # Cheap (a few float unpacks) — fine to run inline on the spin thread.
        self.health.mark_received(config.ODOMETRY_TOPIC)
        p = msg.pose.pose.position
        q = msg.pose.pose.orientation
        self.pointcloud.set_pose(p.x, p.y, p.z, q.x, q.y, q.z, q.w)

    # ── readers (called from asyncio handlers) ───────────────────────────────

    def latest_jpeg(self, cam_id: str) -> Optional[bytes]:
        with self._frame_lock:
            return self._latest_jpeg.get(cam_id)

    def mjpeg_generator(self, cam_id: str, target_hz: float = 12.0):
        """Blocking generator (run in a threadpool) yielding multipart MJPEG chunks."""
        period = 1.0 / target_hz
        boundary = b"--frame"
        while self._running:
            frame = self.latest_jpeg(cam_id)
            if frame is not None:
                yield (
                    boundary + b"\r\nContent-Type: image/jpeg\r\nContent-Length: "
                    + str(len(frame)).encode() + b"\r\n\r\n" + frame + b"\r\n"
                )
            time.sleep(period)