import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import os
import time
from rcl_interfaces.msg import SetParametersResult

from camera_info_manager import CameraInfoManager
from sensor_msgs.msg import CameraInfo

DEFAULT_CAMERA_WIDTH = 1280
DEFAULT_CAMERA_HEIGHT = 480
DEFAULT_TARGET_FPS = 30
DIAG_INTERVAL_SEC = 1.0
 
class CameraNode(Node):
    def __init__(self):
        super().__init__('camera_node')
        left_camera_info_url = 'file:///home/shen/openarm_test/calibrationdata/left.yaml'
        right_camera_info_url = 'file:///home/shen/openarm_test/calibrationdata/right.yaml'
        self.left_camera_info_manager = CameraInfoManager(self, cname='left', url=left_camera_info_url, namespace='/stereo_camera/left')
        self.right_camera_info_manager = CameraInfoManager(self, cname='right', url=right_camera_info_url, namespace='/stereo_camera/right')

        # 设置相机信息文件的路径（不带.yaml后缀）
        # 假设你的yaml文件放在 /home/username/camera_info/ 目录下，名为 left.yaml 和 right.yaml


        # self.left_camera_info_manager.setCameraInfoURL(left_camera_info_url)
        # self.right_camera_info_manager.setCameraInfoURL(right_camera_info_url)
        # self.left_camera_info_manager.setURL(left_camera_info_url)
        # self.right_camera_info_manager.setURL(right_camera_info_url)

        
        self.left_camera_info_manager.loadCameraInfo()
        self.right_camera_info_manager.loadCameraInfo()
        # 然后，你可以通过CameraInfoManager获取CameraInfo消息
        self.left_camera_info = self.left_camera_info_manager.getCameraInfo()
        self.right_camera_info = self.right_camera_info_manager.getCameraInfo()

        # 注意：你需要确保yaml文件已经正确放置在指定路径，并且格式正确。
        self.left_camera_info_pub = self.create_publisher(CameraInfo, '/stereo_camera/left/camera_info', 10)
        self.right_camera_info_pub = self.create_publisher(CameraInfo, '/stereo_camera/right/camera_info', 10)

        self.bridge = CvBridge()
        self.publisher_left = self.create_publisher(Image, '/camera/left/image_raw', 10)
        self.publisher_right = self.create_publisher(Image, '/camera/right/image_raw', 10)
        self.publisher_raw = self.create_publisher(Image, '/image_1', 10)

        # 输入源参数：支持 usb/rtsp 两种来源，并支持运行时切换。
        self.declare_parameter('source_type', 'usb')
        self.declare_parameter('usb_index', 0)
        self.declare_parameter('rtsp_url', '')
        self.declare_parameter('rtsp_transport', 'tcp')
        self.declare_parameter('camera_width', DEFAULT_CAMERA_WIDTH)
        self.declare_parameter('camera_height', DEFAULT_CAMERA_HEIGHT)
        self.declare_parameter('target_fps', DEFAULT_TARGET_FPS)
        self.declare_parameter('usb_use_mjpg', True)
        self.declare_parameter('split_stereo', False)

        self.add_on_set_parameters_callback(self._on_set_parameters)

        self.source_type = self.get_parameter('source_type').value
        self.usb_index = int(self.get_parameter('usb_index').value)
        self.rtsp_url = self.get_parameter('rtsp_url').value
        self.rtsp_transport = self.get_parameter('rtsp_transport').value
        self.camera_width = int(self.get_parameter('camera_width').value)
        self.camera_height = int(self.get_parameter('camera_height').value)
        self.target_fps = int(self.get_parameter('target_fps').value)
        self.usb_use_mjpg = bool(self.get_parameter('usb_use_mjpg').value)
        self.split_stereo = bool(self.get_parameter('split_stereo').value)

        self.diag_last_time = time.perf_counter()
        self.capture_ok_count = 0
        self.publish_count = 0
        self.read_cost_sum = 0.0
        self.convert_cost_sum = 0.0
        self.publish_cost_sum = 0.0
        self.total_cost_sum = 0.0
        self.failed_read_count = 0
        self.reopen_requested = False
        self.reopen_reason = ''
        self.last_reopen_try = 0.0

        self.get_logger().info('正在初始化摄像头...')

        self.cap = None
        self._open_capture()

        timer_period = max(1.0 / max(self.target_fps, 1), 0.001)
        self.timer = self.create_timer(timer_period, self.timer_callback)
        self.get_logger().info('摄像头节点已启动')

    def _on_set_parameters(self, params):
        changed = False
        for p in params:
            if p.name == 'source_type' and p.value not in ('usb', 'rtsp'):
                return SetParametersResult(successful=False, reason='source_type 必须是 usb 或 rtsp')
            if p.name == 'target_fps' and int(p.value) <= 0:
                return SetParametersResult(successful=False, reason='target_fps 必须 > 0')
            if p.name == 'camera_width' and int(p.value) <= 0:
                return SetParametersResult(successful=False, reason='camera_width 必须 > 0')
            if p.name == 'camera_height' and int(p.value) <= 0:
                return SetParametersResult(successful=False, reason='camera_height 必须 > 0')
            if p.name in {
                'source_type', 'usb_index', 'rtsp_url', 'rtsp_transport',
                'camera_width', 'camera_height', 'target_fps', 'usb_use_mjpg',
                'split_stereo'
            }:
                changed = True

        if changed:
            self.reopen_requested = True
            self.reopen_reason = '参数更新触发重连'
        return SetParametersResult(successful=True)

    def _refresh_params(self):
        self.source_type = self.get_parameter('source_type').value
        self.usb_index = int(self.get_parameter('usb_index').value)
        self.rtsp_url = self.get_parameter('rtsp_url').value
        self.rtsp_transport = self.get_parameter('rtsp_transport').value
        self.camera_width = int(self.get_parameter('camera_width').value)
        self.camera_height = int(self.get_parameter('camera_height').value)
        self.target_fps = int(self.get_parameter('target_fps').value)
        self.usb_use_mjpg = bool(self.get_parameter('usb_use_mjpg').value)
        self.split_stereo = bool(self.get_parameter('split_stereo').value)

    def _open_capture(self):
        self._refresh_params()

        if self.cap is not None:
            self.cap.release()
            self.cap = None

        if self.source_type == 'usb':
            self.get_logger().info(f'正在打开 USB 摄像头 /dev/video{self.usb_index} ...')
            cap = cv2.VideoCapture(self.usb_index, cv2.CAP_V4L2)
            if not cap.isOpened():
                raise RuntimeError(f'无法打开 USB 摄像头 /dev/video{self.usb_index}')

            if self.usb_use_mjpg:
                cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
            cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.camera_width)
            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.camera_height)
            cap.set(cv2.CAP_PROP_FPS, self.target_fps)
            cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            self.cap = cap
        else:
            if not self.rtsp_url:
                raise RuntimeError('source_type=rtsp 时必须设置 rtsp_url')
            transport = str(self.rtsp_transport).lower()
            if transport in ('tcp', 'udp'):
                os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = f'rtsp_transport;{transport}'
            self.get_logger().info(f'正在打开 RTSP 流: {self.rtsp_url}')
            cap = cv2.VideoCapture(self.rtsp_url, cv2.CAP_FFMPEG)
            if not cap.isOpened():
                cap = cv2.VideoCapture(self.rtsp_url)
            if not cap.isOpened():
                raise RuntimeError('无法打开 RTSP 流，请检查 url、网络或认证信息')

            cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            self.cap = cap

        # 预热并读取一帧，确保数据流有效。
        for _ in range(3):
            self.cap.read()
        ret, frame = self.cap.read()
        if not ret or frame is None:
            raise RuntimeError('视频源已打开但无法读取图像帧')

        real_w = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        real_h = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        real_fps = self.cap.get(cv2.CAP_PROP_FPS)
        fourcc = int(self.cap.get(cv2.CAP_PROP_FOURCC))
        fourcc_str = ''.join([chr((fourcc >> 8 * j) & 0xFF) for j in range(4)])
        self.get_logger().info(
            f'视频源打开成功: source={self.source_type}, {real_w}x{real_h}, fps={real_fps:.2f}, fourcc={fourcc_str}'
        )

        self.failed_read_count = 0
        self.diag_last_time = time.perf_counter()
        self.capture_ok_count = 0
        self.publish_count = 0
        self.read_cost_sum = 0.0
        self.convert_cost_sum = 0.0
        self.publish_cost_sum = 0.0
        self.total_cost_sum = 0.0

    def _try_reopen_if_needed(self):
        if not self.reopen_requested:
            return

        now = time.perf_counter()
        if now - self.last_reopen_try < 1.0:
            return

        self.last_reopen_try = now
        try:
            self.get_logger().warn(f'尝试重连视频源: {self.reopen_reason}')
            self._open_capture()
            self.reopen_requested = False
            self.reopen_reason = ''
            self.get_logger().info('重连视频源成功')
        except Exception as e:
            self.get_logger().error(f'重连失败: {e}')
    
    def timer_callback(self):
        self._try_reopen_if_needed()
        if self.cap is None or not self.cap.isOpened():
            return

        t0 = time.perf_counter()
        ret, frame = self.cap.read()
        t1 = time.perf_counter()

        if not ret or frame is None:
            self.failed_read_count += 1
            if self.failed_read_count >= 30:
                self.reopen_requested = True
                self.reopen_reason = '连续读取失败'
            return

        self.failed_read_count = 0

        stamp = self.get_clock().now().to_msg()
        if self.split_stereo:
            frame_width = frame.shape[1]
            half_width = frame_width // 2
            msg_left = self.bridge.cv2_to_imgmsg(frame[:, :half_width], encoding='bgr8')
            msg_right = self.bridge.cv2_to_imgmsg(frame[:, half_width:], encoding='bgr8')
            msg_raw= self.bridge.cv2_to_imgmsg(frame, encoding='bgr8')
        else:
            msg_left = None
            msg_right = None
            msg_raw = self.bridge.cv2_to_imgmsg(frame, encoding='bgr8')
        t2 = time.perf_counter()

        if msg_left is not None and msg_right is not None:
            msg_left.header.stamp = stamp
            msg_right.header.stamp = stamp
            msg_left.header.frame_id = 'camera_frame_left'
            msg_right.header.frame_id = 'camera_frame_right'

        if msg_raw is not None:
            msg_raw.header.stamp = stamp
            msg_raw.header.frame_id = 'camera_frame_raw'

        if msg_left is not None and msg_right is not None:
            self.publisher_left.publish(msg_left)
            self.publisher_right.publish(msg_right)
        if msg_raw is not None:
            self.publisher_raw.publish(msg_raw)
        t3 = time.perf_counter()

        # 发布相机信息
        self.left_camera_info_pub.publish(self.left_camera_info)
        self.right_camera_info_pub.publish(self.right_camera_info)
        t4 = time.perf_counter()

        self.capture_ok_count += 1
        self.publish_count += 1
        self.read_cost_sum += (t1 - t0)
        self.convert_cost_sum += (t2 - t1)
        self.publish_cost_sum += (t3 - t2)
        self.total_cost_sum += (t4 - t0)

        now = t4
        elapsed = now - self.diag_last_time
        if elapsed >= DIAG_INTERVAL_SEC:
            cap_fps = self.capture_ok_count / elapsed
            pub_fps = self.publish_count / elapsed
            avg_read_ms = (self.read_cost_sum / max(self.capture_ok_count, 1)) * 1000.0
            avg_convert_ms = (self.convert_cost_sum / max(self.publish_count, 1)) * 1000.0
            avg_publish_ms = (self.publish_cost_sum / max(self.publish_count, 1)) * 1000.0
            avg_total_ms = (self.total_cost_sum / max(self.publish_count, 1)) * 1000.0

            self.get_logger().info(
                f'[diag] source={self.source_type}, cap_fps={cap_fps:.2f}, pub_fps={pub_fps:.2f}, '
                f'read={avg_read_ms:.2f}ms, convert={avg_convert_ms:.2f}ms, '
                f'publish={avg_publish_ms:.2f}ms, total={avg_total_ms:.2f}ms'
            )

            self.diag_last_time = now
            self.capture_ok_count = 0
            self.publish_count = 0
            self.read_cost_sum = 0.0
            self.convert_cost_sum = 0.0
            self.publish_cost_sum = 0.0
            self.total_cost_sum = 0.0

    def destroy_node(self):
        if self.cap is not None:
            self.cap.release()
            self.cap = None
        super().destroy_node()

        


def main():
    rclpy.init()
    node = CameraNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()
 
if __name__ == '__main__':
    main()
