"""
ir_slam_bridge_node.py — Bridge D435i IR1 stream to RGB-style topics for SLAM.

When the vision mode is 'night' or 'ir_only', visible-light cameras are
unusable. RTAB-Map (and other SLAM front-ends) expect a colour image. This
node re-publishes the D435i infra1 (mono8) stream as fake RGB (rgb8) so that
RTAB-Map's visual front-end continues to operate without reconfiguration.

Also re-publishes the camera_info with the same frame_id so that the pipeline
stays consistent.

When mode is 'day' or 'twilight' the bridge is inactive (no republish) to
avoid double-feeding RTAB-Map.

Subscribes:
  /saltybot/vision_mode                std_msgs/String        mode gate
  /camera/infra1/image_rect_raw        sensor_msgs/Image      IR source (mono8)
  /camera/infra1/camera_info           sensor_msgs/CameraInfo IR intrinsics

Publishes (when active):
  /saltybot/ir_bridge/image_raw        sensor_msgs/Image      grey-as-rgb8
  /saltybot/ir_bridge/camera_info      sensor_msgs/CameraInfo passthrough

Parameters:
  active_modes    str   'night,ir_only'   (comma-separated)
  output_encoding str   'rgb8'
"""

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy

import numpy as np
import cv2
from cv_bridge import CvBridge

from sensor_msgs.msg import Image, CameraInfo
from std_msgs.msg import String


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)
_LATCHED_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
    durability=DurabilityPolicy.TRANSIENT_LOCAL,
)
_RELIABLE_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=10,
)


class IRSlamBridgeNode(Node):

    def __init__(self):
        super().__init__('ir_slam_bridge')

        self.declare_parameter('active_modes',    'night,ir_only')
        self.declare_parameter('output_encoding', 'rgb8')

        active_str = self.get_parameter('active_modes').value
        self._active_modes: set[str] = {m.strip() for m in active_str.split(',')}
        self._out_enc: str           = self.get_parameter('output_encoding').value

        self._bridge = CvBridge()
        self._active = False
        self._latest_info: CameraInfo | None = None

        # Mode gate
        self.create_subscription(
            String, '/saltybot/vision_mode', self._on_mode, _RELIABLE_QOS
        )

        # IR source
        self.create_subscription(
            Image, '/camera/infra1/image_rect_raw', self._on_ir_image, _SENSOR_QOS
        )
        self.create_subscription(
            CameraInfo, '/camera/infra1/camera_info', self._on_camera_info, _LATCHED_QOS
        )

        # Bridge output
        self._img_pub  = self.create_publisher(
            Image,      '/saltybot/ir_bridge/image_raw',   10
        )
        self._info_pub = self.create_publisher(
            CameraInfo, '/saltybot/ir_bridge/camera_info', 10
        )

        self.get_logger().info(
            f'ir_slam_bridge ready — active in modes: {sorted(self._active_modes)}'
        )

    # ── Subscriptions ─────────────────────────────────────────────────────────

    def _on_mode(self, msg: String) -> None:
        was_active = self._active
        self._active = msg.data in self._active_modes
        if self._active != was_active:
            state = 'ACTIVE' if self._active else 'idle'
            self.get_logger().info(
                f'[ir_bridge] mode={msg.data} → bridge {state}'
            )

    def _on_camera_info(self, msg: CameraInfo) -> None:
        self._latest_info = msg

    def _on_ir_image(self, msg: Image) -> None:
        if not self._active:
            return

        try:
            mono = self._bridge.imgmsg_to_cv2(msg, desired_encoding='mono8')
        except Exception as exc:
            self.get_logger().warning(f'[ir_bridge] decode error: {exc}')
            return

        # Convert mono8 → target encoding
        if self._out_enc == 'rgb8':
            colour = cv2.cvtColor(mono, cv2.COLOR_GRAY2RGB)
        elif self._out_enc == 'bgr8':
            colour = cv2.cvtColor(mono, cv2.COLOR_GRAY2BGR)
        else:
            colour = mono   # passthrough

        try:
            out_msg = self._bridge.cv2_to_imgmsg(colour, encoding=self._out_enc)
        except Exception as exc:
            self.get_logger().warning(f'[ir_bridge] encode error: {exc}')
            return

        out_msg.header = msg.header
        self._img_pub.publish(out_msg)

        if self._latest_info is not None:
            self._latest_info.header = msg.header
            self._info_pub.publish(self._latest_info)


def main(args=None):
    rclpy.init(args=args)
    node = IRSlamBridgeNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
