import cv2
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import QLabel, QVBoxLayout, QWidget


class VideoDisplayer(QWidget):
    def __init__(self, frame_getter=None) -> None:
        super().__init__()
        self.frame_getter = frame_getter
        self.frame_label = QLabel()
        self.frame_label.setAlignment(Qt.AlignCenter)
        self.layout = QVBoxLayout()
        self.layout.addWidget(self.frame_label)
        self.setLayout(self.layout)

        self.timer = QTimer(self)
        self.timer.timeout.connect(self._update_frame)

    def set_frame_generator(self, frame_getter):
        self.frame_getter = frame_getter

    def start(self):
        self.timer.start(1)  # 每 1 毫秒获取一次帧

    def pause(self):
        self.timer.stop()

    def resume(self):
        self.timer.start(1)

    def close(self):
        self.timer.stop()

    def _update_frame(self):
        if self.frame_getter is None:
            return

        _frame = self.frame_getter()
        if _frame is None:
            return

        try:
            frame = cv2.cvtColor(_frame, cv2.COLOR_BGR2RGB)
            height, width, channel = frame.shape
            step = channel * width
            q_img = QImage(frame.data, width, height, step, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(q_img)
            scaled_pixmap = pixmap.scaled(self.frame_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
            self.frame_label.setPixmap(scaled_pixmap)
        except Exception as e:
            print(f"Error processing frame: {e}")