import rclpy
from sensor_msgs.msg import CompressedImage, Image
from cv_bridge import CvBridge
from rclpy.node import Node
import cv2


class MultiImageDecompressor(Node):
    count = 0

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

        self.bridge = CvBridge()
        self.image_publishers = {}

        # 订阅多个CompressedImage话题
        topics_to_subscribe = ['/usb_cam_front/image_raw/compressed', '/usb_cam_right/image_raw/compressed',
                               '/usb_cam_rear/image_raw/compressed', '/usb_cam_left/image_raw/compressed',
                               '/usb_cam_top/image_raw/compressed', "/fisheye1", "/fisheye2", "/fisheye3", "/fisheye4"]
        for topic in topics_to_subscribe:
            if "compressed" in topic:
                self.image_publishers[topic] = self.create_publisher(
                    Image,
                    topic.replace('/compressed', '/image'),  # 更改话题名称以适应Image
                    10
                )
            else:
                self.image_publishers[topic] = self.create_publisher(
                    Image,
                    topic + '_image',  # 更改话题名称以适应Image
                    10
                )

            self.create_subscription(
                CompressedImage,
                topic,
                lambda msg, topic=topic: self.compressed_image_callback(msg, topic),
                10
            )

    def compressed_image_callback(self, msg, topic):
        print(f"topic is {topic}")
        # 解压缩图像
        try:
            cv_image = self.bridge.compressed_imgmsg_to_cv2(msg, desired_encoding="bgr8")
        except Exception as e:
            self.get_logger().error(f"Error converting compressed image to cv::Mat: {str(e)}")
            return
        #
        # cv2.imwrite(f"/root/ws/data/cam_top/frame_{msg.header.frame_id}_{MultiImageDecompressor.count}.jpg", cv_image)
        # MultiImageDecompressor.count += 1
        # # 创建新的Image消息
        image_msg = self.bridge.cv2_to_imgmsg(cv_image, encoding="bgr8")
        image_msg.header = msg.header

        # 发布解压缩后的Image消息
        self.image_publishers[topic].publish(image_msg)


def main(args=None):
    rclpy.init(args=args)
    multi_image_decompressor = MultiImageDecompressor()
    rclpy.spin(multi_image_decompressor)
    multi_image_decompressor.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
