grpc_client/ui/camera_page.py

import time
from datetime import datetime
from pathlib import Path
from typing import Callable, Optional

from PyQt5 import QtCore, QtGui, QtWidgets

try:
    import av
except Exception as exc:
    av = None
    AV_IMPORT_ERROR = exc
else:
    AV_IMPORT_ERROR = None

try:
    import cv2
except Exception as exc:
    cv2 = None
    CV2_IMPORT_ERROR = exc
else:
    CV2_IMPORT_ERROR = None

from clients._path_setup import ensure_paths

ensure_paths()

from clients.camera_client import CameraClient


class CameraStreamThread(QtCore.QThread):
    frame_received = QtCore.pyqtSignal(object, int, int)
    status_update = QtCore.pyqtSignal(dict)
    error_occurred = QtCore.pyqtSignal(str)
    keyframe_received = QtCore.pyqtSignal()

    def __init__(self, client, device_id, codec="h264", wait_for_keyframe=True):
        super().__init__()
        self.client = client
        self.device_id = device_id
        self.codec_name = codec
        self.wait_for_keyframe = wait_for_keyframe
        self.streaming = False
        self.decoder = None
        self.decoder_error = None
        self.frame_count = 0
        self.fps_counter = 0
        self.fps_time = time.time()
        self.keyframe_count = 0
        self._init_decoder()

    def _init_decoder(self):
        if av is None:
            self.decoder = None
            self.decoder_error = f"PyAV unavailable: {AV_IMPORT_ERROR}"
            return
        try:
            self.decoder = av.CodecContext.create(self.codec_name, "r")
            self.decoder.thread_count = 1
            if hasattr(self.decoder, "flags"):
                try:
                    self.decoder.flags |= av.codec.context.AV_CODEC_FLAG_LOW_DELAY
                except Exception:
                    pass
        except Exception as exc:
            self.decoder = None
            self.decoder_error = f"Decoder init failed: {exc}"

    def _decode(self, frame_data):
        if self.decoder is None or not frame_data:
            return None

        packet = av.Packet(frame_data)
        frames = self.decoder.decode(packet)
        if not frames:
            return None

        return frames[0].to_ndarray(format="rgb24")

    def _start_camera(self):
        response = self.client.start_camera(self.device_id)
        if not response.header.success:
            self.error_occurred.emit(response.header.error_message or "StartCamera failed")
            return False
        return True

    def _stop_camera(self):
        try:
            self.client.stop_camera(self.device_id)
        except Exception:
            pass

    def run(self):
        if self.decoder is None:
            if self.decoder_error:
                self.error_occurred.emit(self.decoder_error)
            return

        try:
            if not self._start_camera():
                return

            self.streaming = True
            self.frame_count = 0
            self.fps_counter = 0
            self.fps_time = time.time()
            self.keyframe_count = 0
            keyframe_seen = False

            responses = self.client.get_rgb_image_stream(
                device_id=self.device_id,
                stop_requested=lambda: not self.streaming,
            )

            for response in responses:
                if not self.streaming:
                    break

                if not response.header.success:
                    raise RuntimeError(response.header.error_message or "GetRGBImageStream failed")

                frame_data = response.color_frame.data
                is_keyframe = response.color_frame.is_key_frame

                if self.wait_for_keyframe and not keyframe_seen and not is_keyframe:
                    continue

                if is_keyframe and not keyframe_seen:
                    keyframe_seen = True
                    self.keyframe_count += 1
                    self.keyframe_received.emit()

                image = self._decode(frame_data)
                if image is None:
                    continue

                self.fps_counter += 1
                self.frame_count += 1

                now = time.time()
                if now - self.fps_time >= 1.0:
                    self.status_update.emit(
                        {
                            "fps": self.fps_counter,
                            "frames": self.frame_count,
                            "keyframes": self.keyframe_count,
                        }
                    )
                    self.fps_counter = 0
                    self.fps_time = now

                height, width = image.shape[:2]
                self.frame_received.emit(image, width, height)
        except Exception as exc:
            self.error_occurred.emit(f"Stream error: {exc}")
        finally:
            self.streaming = False
            self._stop_camera()

    def stop(self):
        self.streaming = False


class CameraPage(QtWidgets.QWidget):
    CODEC_OPTIONS = {
        "H.264 (AVC)": "h264",
        "H.265 (HEVC)": "hevc",
    }

    def __init__(
        self,
        address_provider: Optional[Callable[[], str]] = None,
        default_device_id: str = "right_hand_cam",
        parent=None,
    ):
        super().__init__(parent)
        self.address_provider = address_provider
        self.connected = False
        self.streaming = False
        self.recording = False
        self.stream_client = None
        self.stream_thread = None
        self.video_writer = None
        self.current_frame = None
        self.frame_width = 0
        self.frame_height = 0
        self.record_path = Path("./recordings")
        self.record_path.mkdir(exist_ok=True)

        self._create_ui(default_device_id)

        self.display_timer = QtCore.QTimer(self)
        self.display_timer.timeout.connect(self._update_display)
        self.display_timer.start(30)

    def _create_ui(self, default_device_id: str):
        layout = QtWidgets.QVBoxLayout(self)

        conn_group = QtWidgets.QGroupBox("Camera")
        conn_layout = QtWidgets.QGridLayout(conn_group)

        conn_layout.addWidget(QtWidgets.QLabel("Device ID:"), 0, 0)
        self.device_id = QtWidgets.QLineEdit(default_device_id)
        conn_layout.addWidget(self.device_id, 0, 1)

        conn_layout.addWidget(QtWidgets.QLabel("Codec:"), 0, 2)
        self.codec_combo = QtWidgets.QComboBox()
        self.codec_combo.addItems(list(self.CODEC_OPTIONS))
        self.codec_combo.setCurrentText("H.264 (AVC)")
        conn_layout.addWidget(self.codec_combo, 0, 3)

        self.wait_keyframe = QtWidgets.QCheckBox("Wait for first keyframe")
        self.wait_keyframe.setChecked(True)
        conn_layout.addWidget(self.wait_keyframe, 0, 4, 1, 2)

        if self.address_provider is None:
            conn_layout.addWidget(QtWidgets.QLabel("Server IP:"), 1, 0)
            self.server_ip = QtWidgets.QLineEdit("192.168.0.222")
            conn_layout.addWidget(self.server_ip, 1, 1)

            conn_layout.addWidget(QtWidgets.QLabel("Port:"), 1, 2)
            self.server_port = QtWidgets.QLineEdit("50052")
            self.server_port.setMaximumWidth(100)
            conn_layout.addWidget(self.server_port, 1, 3)
        else:
            conn_layout.addWidget(QtWidgets.QLabel("Server:"), 1, 0)
            self.address_label = QtWidgets.QLabel("")
            self.address_label.setStyleSheet("font-weight: bold;")
            conn_layout.addWidget(self.address_label, 1, 1, 1, 3)

        conn_layout.addWidget(QtWidgets.QLabel("Timeout (s):"), 1, 4)
        self.timeout_spin = QtWidgets.QSpinBox()
        self.timeout_spin.setRange(1, 10)
        self.timeout_spin.setValue(3)
        conn_layout.addWidget(self.timeout_spin, 1, 5)

        conn_layout.addWidget(QtWidgets.QLabel("Record Dir:"), 2, 0)
        self.record_dir = QtWidgets.QLineEdit(str(self.record_path))
        conn_layout.addWidget(self.record_dir, 2, 1, 1, 4)

        self.browse_btn = QtWidgets.QPushButton("Browse")
        self.browse_btn.clicked.connect(self._browse_dir)
        conn_layout.addWidget(self.browse_btn, 2, 5)

        button_layout = QtWidgets.QHBoxLayout()
        self.connect_btn = QtWidgets.QPushButton("Connect")
        self.connect_btn.clicked.connect(self._connect)
        button_layout.addWidget(self.connect_btn)

        self.preview_btn = QtWidgets.QPushButton("Preview")
        self.preview_btn.clicked.connect(self._toggle_preview)
        self.preview_btn.setEnabled(False)
        button_layout.addWidget(self.preview_btn)

        self.record_btn = QtWidgets.QPushButton("Record")
        self.record_btn.clicked.connect(self._toggle_recording)
        self.record_btn.setEnabled(False)
        button_layout.addWidget(self.record_btn)

        self.disconnect_btn = QtWidgets.QPushButton("Disconnect")
        self.disconnect_btn.clicked.connect(self._disconnect)
        self.disconnect_btn.setEnabled(False)
        button_layout.addWidget(self.disconnect_btn)
        button_layout.addStretch()
        conn_layout.addLayout(button_layout, 3, 0, 1, 6)

        layout.addWidget(conn_group)

        status_group = QtWidgets.QGroupBox("Stream Status")
        status_layout = QtWidgets.QHBoxLayout(status_group)

        self.conn_status = QtWidgets.QLabel("Disconnected")
        self.conn_status.setStyleSheet("color: red;")
        status_layout.addWidget(self.conn_status)

        status_layout.addWidget(QtWidgets.QLabel("|"))
        self.codec_status = QtWidgets.QLabel("Codec: H.264 (AVC)")
        status_layout.addWidget(self.codec_status)

        status_layout.addWidget(QtWidgets.QLabel("|"))
        self.fps_status = QtWidgets.QLabel("FPS: 0")
        status_layout.addWidget(self.fps_status)

        status_layout.addWidget(QtWidgets.QLabel("|"))
        self.frame_status = QtWidgets.QLabel("Frames: 0")
        status_layout.addWidget(self.frame_status)

        status_layout.addWidget(QtWidgets.QLabel("|"))
        self.resolution_status = QtWidgets.QLabel("Resolution: -")
        status_layout.addWidget(self.resolution_status)

        status_layout.addWidget(QtWidgets.QLabel("|"))
        self.keyframe_status = QtWidgets.QLabel("Keyframe: waiting")
        self.keyframe_status.setStyleSheet("color: orange;")
        status_layout.addWidget(self.keyframe_status)

        status_layout.addWidget(QtWidgets.QLabel("|"))
        self.record_status = QtWidgets.QLabel("Record: idle")
        self.record_status.setStyleSheet("color: gray;")
        status_layout.addWidget(self.record_status)
        status_layout.addStretch()

        layout.addWidget(status_group)

        video_group = QtWidgets.QGroupBox("Preview")
        video_layout = QtWidgets.QVBoxLayout(video_group)

        self.video_label = QtWidgets.QLabel()
        self.video_label.setAlignment(QtCore.Qt.AlignCenter)
        self.video_label.setMinimumSize(640, 480)
        self.video_label.setStyleSheet("background-color: black; border: 1px solid gray;")
        video_layout.addWidget(self.video_label)

        layout.addWidget(video_group, 1)

        info_frame = QtWidgets.QFrame()
        info_frame.setFrameStyle(QtWidgets.QFrame.StyledPanel | QtWidgets.QFrame.Sunken)
        info_layout = QtWidgets.QHBoxLayout(info_frame)
        info_label = QtWidgets.QLabel(
            "Connect to the camera service, start preview, then optionally record decoded frames."
        )
        info_label.setStyleSheet("color: gray; padding: 5px;")
        info_layout.addWidget(info_label)
        layout.addWidget(info_frame)

        self.status_bar = QtWidgets.QStatusBar()
        self.status_bar.showMessage("Idle")
        layout.addWidget(self.status_bar)

    def refresh_address_hint(self):
        if self.address_provider is None or not hasattr(self, "address_label"):
            return
        try:
            self.address_label.setText(self.address_provider())
        except Exception:
            self.address_label.setText("(app settings unavailable)")

    def _browse_dir(self):
        directory = QtWidgets.QFileDialog.getExistingDirectory(
            self, "Select record directory", self.record_dir.text()
        )
        if directory:
            self.record_dir.setText(directory)
            self.record_path = Path(directory)
            self.record_path.mkdir(exist_ok=True)

    def _get_server_addr(self):
        if self.address_provider is not None:
            return self.address_provider()
        return f"{self.server_ip.text().strip()}:{self.server_port.text().strip()}"

    def _get_codec(self):
        return self.CODEC_OPTIONS[self.codec_combo.currentText()]

    def _connect(self):
        if self.connected:
            return

        self.refresh_address_hint()
        self.status_bar.showMessage("Connecting to camera service...")
        self.codec_status.setText(f"Codec: {self.codec_combo.currentText()}")

        try:
            self.stream_client = CameraClient(
                address=self._get_server_addr(),
                timeout=self.timeout_spin.value(),
            )
        except Exception as exc:
            self.stream_client = None
            self._on_error(f"Connect failed: {exc}")
            return

        self.connected = True
        self.conn_status.setText("Connected")
        self.conn_status.setStyleSheet("color: green;")
        self.connect_btn.setEnabled(False)
        self.preview_btn.setEnabled(True)
        self.disconnect_btn.setEnabled(True)
        self.codec_combo.setEnabled(False)
        self.wait_keyframe.setEnabled(False)
        self.timeout_spin.setEnabled(False)
        if self.address_provider is None:
            self.server_ip.setEnabled(False)
            self.server_port.setEnabled(False)
        self.status_bar.showMessage("Connected to camera service", 3000)

    def _disconnect(self):
        if self.streaming:
            self._stop_preview()

        if self.recording:
            self._stop_recording()

        if self.stream_client:
            self.stream_client.close()
            self.stream_client = None

        self.connected = False
        self.stream_thread = None
        self.conn_status.setText("Disconnected")
        self.conn_status.setStyleSheet("color: red;")
        self.connect_btn.setEnabled(True)
        self.preview_btn.setEnabled(False)
        self.preview_btn.setText("Preview")
        self.record_btn.setEnabled(False)
        self.record_btn.setText("Record")
        self.disconnect_btn.setEnabled(False)
        self.codec_combo.setEnabled(True)
        self.wait_keyframe.setEnabled(True)
        self.timeout_spin.setEnabled(True)
        if self.address_provider is None:
            self.server_ip.setEnabled(True)
            self.server_port.setEnabled(True)
        self._clear_preview()
        self._reset_status_labels()
        self.status_bar.showMessage("Disconnected", 3000)

    def _toggle_preview(self):
        if self.streaming:
            self._stop_preview()
        else:
            self._start_preview()

    def _start_preview(self):
        if not self.connected or self.stream_client is None or self.streaming:
            return

        self.stream_thread = CameraStreamThread(
            client=self.stream_client,
            device_id=self.device_id.text().strip(),
            codec=self._get_codec(),
            wait_for_keyframe=self.wait_keyframe.isChecked(),
        )
        self.stream_thread.frame_received.connect(self._on_frame_received)
        self.stream_thread.status_update.connect(self._on_status_update)
        self.stream_thread.error_occurred.connect(self._on_error)
        self.stream_thread.keyframe_received.connect(self._on_keyframe_received)
        self.stream_thread.finished.connect(self._on_stream_finished)

        self.streaming = True
        self.preview_btn.setText("Stop")
        self.record_btn.setEnabled(True)

        if self.wait_keyframe.isChecked():
            self.keyframe_status.setText("Keyframe: waiting")
            self.keyframe_status.setStyleSheet("color: orange;")
        else:
            self.keyframe_status.setText("Keyframe: not required")
            self.keyframe_status.setStyleSheet("color: gray;")

        self.status_bar.showMessage("Starting preview...")
        self.stream_thread.start()

    def _stop_preview(self):
        if self.recording:
            self._stop_recording()

        thread = self.stream_thread
        if thread is not None:
            thread.stop()
            thread.wait(2000)
            if self.stream_thread is thread:
                self._on_stream_finished()
        else:
            self._on_stream_finished()

    def _on_stream_finished(self):
        if self.recording:
            self._stop_recording()
        self.streaming = False
        self.stream_thread = None
        self.preview_btn.setText("Preview")
        self.record_btn.setEnabled(False)
        self.keyframe_status.setText("Keyframe: waiting")
        self.keyframe_status.setStyleSheet("color: orange;")
        self.status_bar.showMessage("Preview stopped", 3000)

    def _toggle_recording(self):
        if self.recording:
            self._stop_recording()
        else:
            self._start_recording()

    def _start_recording(self):
        if not self.streaming:
            QtWidgets.QMessageBox.warning(self, "Warning", "Start preview before recording.")
            return
        if cv2 is None:
            QtWidgets.QMessageBox.critical(self, "Error", f"OpenCV unavailable: {CV2_IMPORT_ERROR}")
            return

        self.record_path = Path(self.record_dir.text())
        self.record_path.mkdir(parents=True, exist_ok=True)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        codec_name = self.codec_combo.currentText().replace(" ", "_")
        filename = self.record_path / f"rec_{self.device_id.text()}_{codec_name}_{timestamp}.mp4"

        if self.current_frame is not None:
            height, width = self.current_frame.shape[:2]
        else:
            width, height = 640, 480

        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        self.video_writer = cv2.VideoWriter(str(filename), fourcc, 30.0, (width, height))

        if not self.video_writer.isOpened():
            self.video_writer = None
            QtWidgets.QMessageBox.critical(self, "Error", "Failed to create the output video file.")
            return

        self.recording = True
        self.record_btn.setText("Stop Recording")
        self.record_status.setText("Record: active")
        self.record_status.setStyleSheet("color: red;")
        self.status_bar.showMessage(f"Recording to {filename.name}", 3000)

    def _stop_recording(self):
        if self.video_writer:
            self.video_writer.release()
            self.video_writer = None

        self.recording = False
        self.record_btn.setText("Record")
        self.record_status.setText("Record: idle")
        self.record_status.setStyleSheet("color: gray;")
        self.status_bar.showMessage("Recording stopped", 3000)

    def _on_frame_received(self, frame, width, height):
        self.current_frame = frame
        self.frame_width = width
        self.frame_height = height

        if self.recording and self.video_writer and cv2 is not None:
            self.video_writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))

    def _update_display(self):
        if self.current_frame is None:
            return

        bytes_per_line = 3 * self.frame_width
        qimage = QtGui.QImage(
            self.current_frame.data,
            self.frame_width,
            self.frame_height,
            bytes_per_line,
            QtGui.QImage.Format_RGB888,
        )

        pixmap = QtGui.QPixmap.fromImage(qimage).scaled(
            self.video_label.size(),
            QtCore.Qt.KeepAspectRatio,
            QtCore.Qt.SmoothTransformation,
        )
        self.video_label.setPixmap(pixmap)
        self.resolution_status.setText(f"Resolution: {self.frame_width}x{self.frame_height}")

    def _on_status_update(self, stats):
        self.fps_status.setText(f"FPS: {stats['fps']}")
        self.frame_status.setText(f"Frames: {stats['frames']}")

    def _on_keyframe_received(self):
        if self.wait_keyframe.isChecked():
            self.keyframe_status.setText("Keyframe: received")
            self.keyframe_status.setStyleSheet("color: green;")

    def _on_error(self, message):
        QtWidgets.QMessageBox.critical(self, "Error", message)
        self.status_bar.showMessage(message, 5000)

    def _clear_preview(self):
        self.video_label.clear()
        self.current_frame = None
        self.frame_width = 0
        self.frame_height = 0

    def _reset_status_labels(self):
        self.fps_status.setText("FPS: 0")
        self.frame_status.setText("Frames: 0")
        self.resolution_status.setText("Resolution: -")
        self.keyframe_status.setText("Keyframe: waiting")
        self.keyframe_status.setStyleSheet("color: orange;")
        self.record_status.setText("Record: idle")
        self.record_status.setStyleSheet("color: gray;")

    def shutdown(self):
        if self.streaming:
            self._stop_preview()
        if self.recording:
            self._stop_recording()
        if self.connected:
            self._disconnect()


class CameraWindow(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Camera Stream Player")
        self.resize(1000, 800)
        self.camera_page = CameraPage(parent=self)
        self.setCentralWidget(self.camera_page)

    def closeEvent(self, event):
        self.camera_page.shutdown()
        super().closeEvent(event)