import cv2
from ultralytics import YOLO
from opcua import Client, ua
import dmPython
import pandas as pd
from datetime import datetime, timedelta
import time


class ConveyorBeltMonitor:
    def __init__(self, model_path, plc_config, cameras_config):
        # 初始化 YOLO model
        self.model = YOLO(model_path)

        # PLC 连接 设置
        self.plc_url = plc_config['url']
        self.plc_username = plc_config['username']
        self.plc_password = plc_config['password']
        self.client = None
        self.plc_timeout = 5  # 设置连接和读取超时时间（秒）
        self.max_retries = 3  # 最大重试次数

        # 摄像头 plc node 设置
        self.cameras = {}
        # 每个摄像头状态
        self.camera_states = {}
        
        self.skip_frames = 20

        self.setup_cameras(cameras_config)

    def setup_cameras(self, cameras_config):
        """Initialize multiple camera connections with their PLC nodes"""
        for camera_id, config in cameras_config.items():
            try:
                cap = cv2.VideoCapture(config['rtsp_url'])
                if not cap.isOpened():
                    print(f"打开 {camera_id} 摄像头失败!")
                    continue
                self.cameras[camera_id] = {
                    'capture': cap,
                    'frame_count': 0,
                    'rtsp_url': config['rtsp_url'],
                    'plc_nodes': config['plc_nodes'],  # Now accepts multiple PLC nodes
                    'failed_read_count': 0
                }
                self.camera_states[camera_id] = {
                    'last_plc_values': {node: None for node in config['plc_nodes']},  # Track last value for each node
                    'has_queried_database': False,
                    'coal_detection_count': 0  # 记录识别到煤的次数
                }
            except Exception as e:
                print(f"Error setting up camera {camera_id}: {e}")

    def connect_plc(self):
        """建立 PLC 连接"""
        self.client = Client(self.plc_url)
        self.client.set_user(self.plc_username)
        self.client.set_password(self.plc_password)
        self.client.timeout = self.plc_timeout
        try:
            self.client.connect()
            print("连接到 OPC UA 服务器")
        except Exception as e:
            print(f"Failed to connect to PLC: {e}")
            self.reconnect_plc()

    def reconnect_plc(self):
        """尝试重新连接 PLC"""
        print("尝试重新连接 PLC...")
        time.sleep(5)  # 等待 5 秒后重试
        self.connect_plc()

    def reconnect_camera(self, camera_id):
        """Attempt to reconnect the camera after failing to read frames"""
        config = self.cameras[camera_id]
        print(f"重连摄像头 {camera_id}...")
        config['capture'].release()  # Release the current capture object
        try:
            cap = cv2.VideoCapture(config['rtsp_url'])
            if cap.isOpened():
                print(f"重连摄像头 {camera_id} 成功")
                config['capture'] = cap
                config['failed_read_count'] = 0  # Reset the failed count
            else:
                print(f"Failed to reconnect to camera {camera_id}")
        except Exception as e:
            print(f"Error reconnecting to camera {camera_id}: {e}")

    def fetch_alarm_data(self):
        """查询 达梦 报警信息"""
        db_config = {
            'host': "10.248.56.208",
            'port': "5236",
            'user': "XMSW",
            'password': "7m*d7Zf04#G%4je"
        }

        try:
            connection = dmPython.Connection(
                user=db_config['user'],
                password=db_config['password'],
                server=db_config['host'],
                port=db_config['port']
            )
            cursor = connection.cursor()
            now = datetime.now()
            two_minutes_ago = now - timedelta(minutes=2)

            query = f"""
                SELECT TAG_NAME, TAG_COMMENT, ALARM_TIME, EVENT_TIME, ALARM_TEXT
                FROM XMSW.T_ALARM
                WHERE ALARM_TIME BETWEEN '{two_minutes_ago}' AND '{now}'
                ORDER BY ALARM_TIME DESC
            """

            cursor.execute(query)
            rows = cursor.fetchall()
            columns = ['TAG_NAME', 'TAG_COMMENT', 'ALARM_TIME', 'EVENT_TIME', 'ALARM_TEXT']
            df = pd.DataFrame(rows, columns=columns)

            cursor.close()
            connection.close()
            return df
        except Exception as e:
            print(f"Failed to fetch alarm data: {e}")
            return pd.DataFrame()  # 返回空 DataFrame

    def process_frame(self, frame):
        """帧 缩放"""
        try:
            dim = (int(frame.shape[1] * 0.5), int(frame.shape[0] * 0.5))
            resized_frame = cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)

            results = self.model(resized_frame)
            annotated_frame = results[0].plot()
            detections = results[0].boxes

            return annotated_frame, detections
        except Exception as e:
            print(f"Error processing frame: {e}")
            return frame, None  # 返回原始帧和空的检测结果

    def read_plc_node(self, node_id, camera_id):
        """读取 PLC 节点值，处理超时"""
        for attempt in range(self.max_retries):
            try:
                node = self.client.get_node(node_id)
                value = node.get_value()
                return value
            except (ua.UaError, TimeoutError) as e:
                print(f"Attempt {attempt + 1}: Failed to read PLC node {node_id}, Error: {e}")
                time.sleep(1)  # 等待 1 秒后重试
        print(f"Max retries reached for node {node_id}. Returning last value.")
        return self.camera_states[camera_id]['last_plc_values'][node_id]  # 如果所有重试失败，返回 上次的值

    def check_plc_conditions(self, camera_id):
        """Check PLC conditions for a specific camera"""
        plc_nodes = self.cameras[camera_id]['plc_nodes']
        current_values = {}
        trigger_condition = False

        last_values = self.camera_states[camera_id]['last_plc_values']

        # 获取 plc node 值
        for node_id in plc_nodes:
            current_values[node_id] = self.read_plc_node(node_id, camera_id)
            print(f"Camera {camera_id}: PLC  {node_id} 当前值: {current_values[node_id]}")
            print(f"Camera {camera_id}: PLC node {node_id} 上次值: {last_values[node_id]}")

        # 判断 运行状态
        trigger_condition = any(
            last_values[node_id] == 3 and current_values[node_id] != 3
            for node_id in plc_nodes
        )

        # Update last values
        self.camera_states[camera_id]['last_plc_values'] = current_values

        return trigger_condition

    def handle_detections(self, camera_id, detections):
        """Handle detection results and PLC communication for specific camera"""
        if detections is not None and detections.shape[0] > 0:
            print(f"Camera {camera_id}: Objects detected: {len(detections)}")

            self.camera_states[camera_id]['coal_detection_count'] += 1
            print(
                f"Camera {camera_id}: Coal detected. Count: {self.camera_states[camera_id]['coal_detection_count']}")

            # 如果识别到煤的次数达到 10 次，查询数据库并保存到 Excel
            if self.camera_states[camera_id]['coal_detection_count'] >= 10:
                current_time = datetime.now().strftime("%Y-%m-%d %H-%M-%S")
                alarm_data = self.fetch_alarm_data()

                excel_file_name = f"胶带机_{camera_id}_{current_time}.xlsx"
                try:
                    alarm_data.to_excel(excel_file_name, index=False, sheet_name='QueryResults', engine='openpyxl')
                    print(f"Camera {camera_id}: Alarm data written to {excel_file_name}")
                    self.camera_states[camera_id]['coal_detection_count'] = 0  # 重置计数
                except Exception as e:
                    print(f"Failed to write alarm data to Excel: {e}")
        else:
            print(f"Camera {camera_id}: No objects detected")

    def run(self):
        """Main monitoring loop"""
        try:
            self.connect_plc()

            while True:
                for camera_id, camera_info in self.cameras.items():
                    # cap = camera_info['capture']
                    # success, frame = cap.read()

                    # if not success:
                    #     camera_info['failed_read_count'] += 1
                    #     print(
                    #         f"Failed to read frame from camera {camera_id}, Attempt: {camera_info['failed_read_count']}")
                    #     if camera_info['failed_read_count'] >= 5:
                    #         self.reconnect_camera(camera_id)
                    #     continue

                    # camera_info['frame_count'] += 1
                    # camera_info['failed_read_count'] = 0
                    # if camera_info['frame_count'] % self.skip_frames != 0:
                    #     continue

                    # annotated_frame, detections = self.process_frame(frame)

                    # 检查 PLC 条件
                    trigger_condition = self.check_plc_conditions(camera_id)

                    if trigger_condition:
                        cap = camera_info['capture']
                        success, frame = cap.read()
                        if not success:
                            camera_info['failed_read_count'] += 1
                            print(
                                f"Failed to read frame from camera {camera_id}, Attempt: {camera_info['failed_read_count']}")
                            if camera_info['failed_read_count'] >= 5:
                                self.reconnect_camera(camera_id)
                            continue
                            
                        annotated_frame, detections = self.process_frame(frame)
                        self.handle_detections(camera_id, detections)

                    cv2.imshow(f"Camera {camera_id}", annotated_frame)

                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        finally:
            self.cleanup()

    def cleanup(self):
        """释放资源"""
        if self.client:
            self.client.disconnect()
        for camera_info in self.cameras.values():
            camera_info['capture'].release()
        cv2.destroyAllWindows()


# Configuration
plc_config = {
    'url': "opc.tcp://10.248.56.206:37802",
    'username': "swxmc",
    'password': "swxmc@147!"
}

cameras_config = {
    'cam218': {
        'rtsp_url': 'rtsp://admin:hik12345+@10.248.53.20',
        'plc_nodes': ["ns=2;s=XMSW_Raw_CON218_RunState11"]  # 单 PLC node
    },
    'cam201': {
        'rtsp_url': 'rtsp://admin:hik12345+@10.248.53.96',
        'plc_nodes': [
            "ns=2;s=XMSW_Raw_CON201_RunState11",
            "ns=2;s=XMSW_Raw_CON201_RunState22"
        ]  # 多 PLC nodes
    }
}

# Create and run monitor
monitor = ConveyorBeltMonitor("best.pt", plc_config, cameras_config)
monitor.run()
上述代码存在一个问题，当plc的值从3变为不等于3的时候执行图像识别，但是图像识别只会在值变化的时候触发一次，为了保证图像识别的准确性，防止有物体漏检，修改代码当plc的值变化时，触发图像识别执行多次，并且保证识别物体次数不小于5次，那么判断为有物体存在，执行数据库查询等操作