#!/usr/bin/env python3
#ros2 run camera_test camera_server

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
import numpy as np
from protocol.srv import CameraService
import cv2
import os
import time
from datetime import datetime


# 配置显示环境（关键修改）
os.environ['OPENCV_VIDEOIO_PRIORITY_GTK'] = '1'
os.environ['DISPLAY'] = ':0'  # 使用主显示器

class AICameraDemo(Node):
    def __init__(self):
        super().__init__('ai_camera_demo')
        
        # 创建服务客户端
        self.camera_client = self.create_client(CameraService, '/camera_service')
        while not self.camera_client.wait_for_service(timeout_sec=1.0):
            self.get_logger().info('等待AI相机服务...')
        
        # 图像订阅
        self.image_subscription = self.create_subscription(
            Image, '/image', self.image_callback, 10
        )
        
        # 截图控制参数
        self.last_capture_time = time.time()
        self.start_time = time.time()
        self.capture_interval = 3.0  # 截图间隔(秒)
        self.total_duration = 30.0   # 总运行时间(秒)
        self.save_dir = "captures"   # 截图保存目录
        os.makedirs(self.save_dir, exist_ok=True)
        
        # 初始化OpenCV窗口
        self.window_name = 'AI Camera Feed'
        try:
            cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
            self.gui_enabled = True
        except:
            self.get_logger().error("无法创建OpenCV窗口，将仅保存截图")
            self.gui_enabled = False

    def start_camera(self):
        """启动相机服务"""
        req = CameraService.Request()
        req.command = 9  # START_IMAGE_PUBLISH
        req.width, req.height, req.fps = 640, 480, 30
        future = self.camera_client.call_async(req)
        rclpy.spin_until_future_complete(self, future)
        return future.result().result == 0

    def stop_camera(self):
        """停止相机服务"""
        req = CameraService.Request()
        req.command = 10  # STOP_IMAGE_PUBLISH
        future = self.camera_client.call_async(req)
        rclpy.spin_until_future_complete(self, future)
        return future.result().result == 0

    def image_callback(self, msg):
        """处理图像并定时截图"""
        current_time = time.time()
        elapsed = current_time - self.start_time
        
        # 超过总时长则停止
        if elapsed >= self.total_duration:
            self.get_logger().info("已完成30秒截图任务")
            raise KeyboardInterrupt
        
        try:
            # 转换图像格式
            if msg.encoding in ['rgb8', 'bgr8']:
                img = np.frombuffer(msg.data, dtype=np.uint8).reshape(
                    msg.height, msg.width, 3
                )
                if msg.encoding == 'bgr8':
                    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                
                # 定时截图逻辑
                if current_time - self.last_capture_time >= self.capture_interval:
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    filename = f"{self.save_dir}/capture_{timestamp}.jpg"
                    cv2.imwrite(filename, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
                    self.get_logger().info(f"已保存截图: {filename}")
                    self.last_capture_time = current_time
                
                # 实时显示
                if self.gui_enabled:
                    cv2.imshow(self.window_name, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
                    cv2.waitKey(1)
                    
        except Exception as e:
            self.get_logger().error(f"图像处理错误: {str(e)}")

def main(args=None):
    rclpy.init(args=args)
    demo = AICameraDemo()
    
    try:
        if demo.start_camera():
            demo.get_logger().info("开始截图任务 (每3秒截图，持续30秒)")
            rclpy.spin(demo)
    except KeyboardInterrupt:
        demo.get_logger().info("手动终止")
    finally:
        demo.stop_camera()
        cv2.destroyAllWindows()
        demo.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()