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

# pip install openpyxl -i https://pypi.tuna.tsinghua.edu.cn/simple
# pip install cryptography -i https://pypi.tuna.tsinghua.edu.cn/simple

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
                }
            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}")

    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']

        excel_file = 'plc_data.xlsx'

        try:
            # 尝试读取现有的Excel文件
            df = pd.read_excel(excel_file)
        except FileNotFoundError:
            # 如果文件不存在，创建一个新的DataFrame
            df = pd.DataFrame(columns=['Camera ID', 'Node ID', 'Last Value', 'Current Value', 'Timestamp'])

        # 获取 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 {node_id} 上次值: {last_values[node_id]}")
            print(f"Camera {camera_id}: PLC node {node_id} 当前值: {current_values[node_id]}")
            if(last_values[node_id] == 3 and current_values[node_id] != 3):
                current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                # 创建一个新的行数据
                new_row = {
                    'Camera ID': camera_id,
                    'Node ID': node_id,
                    'Last Value': last_values[node_id],
                    'Current Value': current_values[node_id],
                    'Timestamp': current_time
                }

                # 将新行转换为 DataFrame
                new_row_df = pd.DataFrame([new_row])

                # 使用 pd.concat 将新行追加到 DataFrame 中
                df = pd.concat([df, new_row_df], ignore_index=True)

                # 将更新后的 DataFrame 保存到 Excel 文件中
                df.to_excel(excel_file, index=False)

        # 判断 运行状态
        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"""
        trigger_condition = self.check_plc_conditions(camera_id)
        if detections is not None and detections.shape[0] > 0:
            print(f"Camera {camera_id}: Objects detected: {len(detections)}")

            # trigger_condition = self.check_plc_conditions(camera_id)
            # has_queried_database = self.camera_states[camera_id]['has_queried_database']

            if trigger_condition: # and not has_queried_database:
                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]['has_queried_database'] = True
                except Exception as e:
                    print(f"Failed to write alarm data to Excel: {e}")
            # elif not trigger_condition:
            #     self.camera_states[camera_id]['has_queried_database'] = False
        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)
                    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()

请检查上述代码是否有代码冗余或者不合适的地方。