import cv2
import PyQt5.QtWidgets as QtWidgets
import PyQt5.QtGui as QtGui
import PyQt5.QtCore as QtCore
import sys
import time


class VideoWidget(QtWidgets.QWidget):
    def __init__(self, rtsp_url):
        super().__init__()
        self.rtsp_url = rtsp_url
        self.cap = cv2.VideoCapture(self.rtsp_url)
        self.timer = QtCore.QTimer(self)
        self.timer.timeout.connect(self.update_frame)
        self.timer.start(30)
        self.start_time = time.time()
        self.video_writer = None
        self.new_video_writer()

    def new_video_writer(self):
        if self.video_writer is not None:
            self.video_writer.release()
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can also try 'avc1'
        filename = f"output_{time.strftime('%Y%m%d%H%M%S')}.mp4"
        ret, frame = self.cap.read()
        if ret:
            height, width, _ = frame.shape
            self.video_writer = cv2.VideoWriter(filename, fourcc, 20.0, (width, height))

    def update_frame(self):
        ret, frame = self.cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            height, width, channel = frame.shape
            step = channel * width
            q_img = QtGui.QImage(frame.data, width, height, step, QtGui.QImage.Format_RGB888)
            self.video.setPixmap(QtGui.QPixmap.fromImage(q_img))

            # Write frame to video
            self.video_writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))

            # Check if 1 hour has passed
            if time.time() - self.start_time >= 3600:
                self.start_time = time.time()
                self.new_video_writer()

    def init_ui(self):
        self.video = QtWidgets.QLabel(self)
        layout = QtWidgets.QVBoxLayout()
        layout.addWidget(self.video)
        self.setLayout(layout)
        self.setWindowTitle('RTSP Video Stream')
        self.setGeometry(100, 100, 800, 600)
        self.show()

    def closeEvent(self, event):
        self.cap.release()
        self.video_writer.release()
        self.timer.stop()
        event.accept()


def main():
    rtsp_url = "rtsp://admin:@yjbabo.ddns.net:999/ch0_1.264"
    app = QtWidgets.QApplication(sys.argv)
    video_widget = VideoWidget(rtsp_url)
    video_widget.init_ui()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()