#!/usr/bin/env python3
"""Expose visualization-compatible ROS topics for the web plugin."""

from collections import defaultdict

import rclpy
from rclpy.node import Node
from rclpy.qos import DurabilityPolicy, QoSProfile, ReliabilityPolicy
from sensor_msgs.msg import CameraInfo, Image, PointCloud2

from bcr_bot.srv import GetVisualTopics


IMAGE_TYPE = "sensor_msgs/msg/Image"
CAMERA_INFO_TYPE = "sensor_msgs/msg/CameraInfo"
POINT_CLOUD_TYPE = "sensor_msgs/msg/PointCloud2"

DEPTH_ENCODINGS = {"16uc1", "32fc1", "mono16"}
RGB_ENCODINGS = {
    "rgb8",
    "bgr8",
    "rgba8",
    "bgra8",
    "8uc3",
    "8uc4",
    "r8g8b8",
    "b8g8r8",
}


class VisualTopicClassifier(Node):
    def __init__(self):
        super().__init__("visual_topic_classifier")

        self._image_encodings = {}
        self._image_subscriptions = {}
        self._activity_subscriptions = {}
        self._active_topics = set()
        self._topics_by_type = defaultdict(set)

        self.create_service(
            GetVisualTopics, "get_visual_topics", self._get_visual_topics_callback
        )
        self.create_timer(1.0, self._scan_topics)
        self._scan_topics()

        self.get_logger().info(
            "Visual topic classifier is ready on service '/get_visual_topics'."
        )

    def _scan_topics(self):
        topics_by_type = defaultdict(set)

        for topic_name, topic_types in self.get_topic_names_and_types():
            for topic_type in topic_types:
                topics_by_type[topic_type].add(topic_name)

            if IMAGE_TYPE in topic_types:
                self._subscribe_to_image_encoding(topic_name)
            if CAMERA_INFO_TYPE in topic_types:
                self._subscribe_for_activity(topic_name, CameraInfo)
            if POINT_CLOUD_TYPE in topic_types:
                self._subscribe_for_activity(topic_name, PointCloud2)

        self._topics_by_type = topics_by_type

    def _subscribe_to_image_encoding(self, topic_name):
        if topic_name in self._image_subscriptions:
            return

        qos = QoSProfile(
            depth=1,
            reliability=ReliabilityPolicy.BEST_EFFORT,
            durability=DurabilityPolicy.VOLATILE,
        )
        self._image_subscriptions[topic_name] = self.create_subscription(
            Image,
            topic_name,
            lambda message, name=topic_name: self._cache_image_encoding(name, message),
            qos,
        )

    def _subscribe_for_activity(self, topic_name, message_type):
        if topic_name in self._activity_subscriptions:
            return

        qos = QoSProfile(
            depth=1,
            reliability=ReliabilityPolicy.BEST_EFFORT,
            durability=DurabilityPolicy.VOLATILE,
        )
        self._activity_subscriptions[topic_name] = self.create_subscription(
            message_type,
            topic_name,
            lambda _, name=topic_name: self._mark_topic_active(name),
            qos,
        )

    def _cache_image_encoding(self, topic_name, message):
        self._mark_topic_active(topic_name)
        encoding = message.encoding.lower()
        if self._image_encodings.get(topic_name) != encoding:
            self._image_encodings[topic_name] = encoding
            self.get_logger().info(
                "Detected image encoding '%s' on '%s'." % (encoding, topic_name)
            )

    def _mark_topic_active(self, topic_name):
        self._active_topics.add(topic_name)

    def _get_visual_topics_callback(self, request, response):
        if request.refresh:
            self._scan_topics()

        response.depth_image_topics = self._image_topics_for(DEPTH_ENCODINGS)
        response.rgb_image_topics = self._image_topics_for(RGB_ENCODINGS)
        response.camera_info_topics = self._active_topics_for(CAMERA_INFO_TYPE)
        response.point_cloud_topics = self._active_topics_for(POINT_CLOUD_TYPE)
        return response

    def _image_topics_for(self, expected_encodings):
        return sorted(
            topic_name
            for topic_name in self._topics_by_type[IMAGE_TYPE]
            if self._image_encodings.get(topic_name) in expected_encodings
        )

    def _active_topics_for(self, topic_type):
        return sorted(
            topic_name
            for topic_name in self._topics_by_type[topic_type]
            if topic_name in self._active_topics
        )


def main(args=None):
    rclpy.init(args=args)
    node = VisualTopicClassifier()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == "__main__":
    main()
