from PySide6.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QSlider, QCheckBox, QFileDialog, QSizePolicy, QGridLayout, QFrame, QSpacerItem
from PySide6.QtCore import Signal, Qt
from PySide6.QtGui import QPixmap, QIcon, QFont

class DiagnosticOverlay(QFrame):
    """Semi-transparent overlay for diagnostics."""
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setStyleSheet("""
            background: rgba(0,0,0,0.5);
            border-radius: 8px;
            color: #fff;
            font-family: 'Consolas', 'SF Mono', 'monospace';
            font-size: 13px;
        """)
        # self.setFixedWidth(260)  # Remove fixed width
        self.setFixedHeight(90)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)  # Allow horizontal stretch
        self.setAttribute(Qt.WA_TransparentForMouseEvents)
        layout = QVBoxLayout(self)
        layout.setContentsMargins(12, 8, 12, 8)
        self.model_label = QLabel("Model: -")
        self.device_label = QLabel("Device: -")
        self.stats_label = QLabel("Cars: 0 | Trucks: 0 | Ped: 0 | TLights: 0 | Moto: 0")
        for w in [self.model_label, self.device_label, self.stats_label]:
            w.setStyleSheet("color: #fff;")
            layout.addWidget(w)
        layout.addStretch(1)

    def update_overlay(self, model, device, cars, trucks, peds, tlights, motorcycles):
        self.model_label.setText(f"Model: {model}")
        self.device_label.setText(f"Device: {device}")
        self.stats_label.setText(f"Cars: {cars} | Trucks: {trucks} | Ped: {peds} | TLights: {tlights} | Moto: {motorcycles}")

class VideoDetectionTab(QWidget):
    file_selected = Signal(str)
    play_clicked = Signal()
    pause_clicked = Signal()
    stop_clicked = Signal()
    detection_toggled = Signal(bool)
    screenshot_clicked = Signal()
    seek_changed = Signal(int)
    auto_select_model_device = Signal()  # New signal for auto model/device selection

    def __init__(self):
        super().__init__()
        self.video_loaded = False
        grid = QGridLayout(self)
        grid.setContentsMargins(32, 24, 32, 24)
        grid.setSpacing(0)
        # File select bar (top)
        file_bar = QHBoxLayout()
        self.file_btn = QPushButton()
        self.file_btn.setIcon(QIcon.fromTheme("folder-video"))
        self.file_btn.setText("Select Video")
        self.file_btn.setStyleSheet("padding: 8px 18px; border-radius: 8px; background: #232323; color: #fff;")
        self.file_label = QLabel("No file selected")
        self.file_label.setStyleSheet("color: #bbb; font-size: 13px;")
        self.file_btn.clicked.connect(self._select_file)
        file_bar.addWidget(self.file_btn)
        file_bar.addWidget(self.file_label)
        file_bar.addStretch(1)
        # Video display area (centered, scalable)
        video_frame = QFrame()
        video_frame.setStyleSheet("""
            background: #121212;
            border: 1px solid #424242;
            border-radius: 8px;
        """)
        video_frame.setMinimumSize(640, 360)
        video_frame.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        video_layout = QVBoxLayout(video_frame)
        video_layout.setContentsMargins(0, 0, 0, 0)
        video_layout.setAlignment(Qt.AlignCenter)
        self.video_label = QLabel()
        self.video_label.setAlignment(Qt.AlignCenter)
        self.video_label.setStyleSheet("background: transparent; color: #888; font-size: 18px;")
        self.video_label.setText("No video loaded. Please select a file.")
        self.video_label.setMinimumSize(640, 360)
        self.video_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        video_layout.addWidget(self.video_label)
        # Diagnostic overlay (now below video, not over it)
        self.overlay = DiagnosticOverlay()
        self.overlay.setStyleSheet(self.overlay.styleSheet() + "border: 1px solid #03DAC5;")
        self.overlay.setFixedHeight(90)
        # FPS and Inference badges (below video)
        self.fps_badge = QLabel("FPS: --")
        self.fps_badge.setStyleSheet("background: #27ae60; color: #fff; border-radius: 12px; padding: 4px 24px; font-weight: bold; font-size: 15px;")
        self.fps_badge.setAlignment(Qt.AlignCenter)
        self.inference_badge = QLabel("Inference: -- ms")
        self.inference_badge.setStyleSheet("background: #2980b9; color: #fff; border-radius: 12px; padding: 4px 24px; font-weight: bold; font-size: 15px;")
        self.inference_badge.setAlignment(Qt.AlignCenter)
        # Horizontal layout for overlay and badges
        self.badge_bar = QHBoxLayout()
        self.badge_bar.setContentsMargins(0, 8, 0, 8)
        self.badge_bar.addWidget(self.fps_badge)
        self.badge_bar.addSpacing(12)
        self.badge_bar.addWidget(self.inference_badge)
        self.badge_bar.addSpacing(18)
        self.badge_bar.addWidget(self.overlay)  # Overlay will stretch to fill right side
        self.badge_bar.addStretch(10)
        video_layout.addStretch(1)  # Push badge bar to the bottom
        video_layout.addLayout(self.badge_bar)
        # Control bar (bottom)
        control_bar = QHBoxLayout()
        control_bar.setContentsMargins(0, 16, 0, 0)
        # Playback controls
        self.play_btn = QPushButton()
        self.play_btn.setIcon(QIcon.fromTheme("media-playback-start"))
        self.play_btn.setToolTip("Play")
        self.play_btn.setFixedSize(48, 48)
        self.play_btn.setEnabled(False)
        self.play_btn.setStyleSheet(self._button_style())
        self.pause_btn = QPushButton()
        self.pause_btn.setIcon(QIcon.fromTheme("media-playback-pause"))
        self.pause_btn.setToolTip("Pause")
        self.pause_btn.setFixedSize(48, 48)
        self.pause_btn.setEnabled(False)
        self.pause_btn.setStyleSheet(self._button_style())
        self.stop_btn = QPushButton()
        self.stop_btn.setIcon(QIcon.fromTheme("media-playback-stop"))
        self.stop_btn.setToolTip("Stop")
        self.stop_btn.setFixedSize(48, 48)
        self.stop_btn.setEnabled(False)
        self.stop_btn.setStyleSheet(self._button_style())
        for btn, sig in zip([self.play_btn, self.pause_btn, self.stop_btn], [self.play_clicked.emit, self.pause_clicked.emit, self.stop_clicked.emit]):
            btn.clicked.connect(sig)
        control_bar.addWidget(self.play_btn)
        control_bar.addWidget(self.pause_btn)
        control_bar.addWidget(self.stop_btn)
        control_bar.addSpacing(16)
        # Progress bar
        self.progress = QSlider(Qt.Horizontal)
        self.progress.setStyleSheet("QSlider::groove:horizontal { height: 6px; background: #232323; border-radius: 3px; } QSlider::handle:horizontal { background: #03DAC5; border-radius: 8px; width: 18px; }")
        self.progress.setMinimumWidth(240)
        self.progress.setEnabled(False)
        self.progress.valueChanged.connect(self.seek_changed.emit)
        control_bar.addWidget(self.progress, 2)
        self.timestamp = QLabel("00:00 / 00:00")
        self.timestamp.setStyleSheet("color: #bbb; font-size: 13px;")
        control_bar.addWidget(self.timestamp)
        control_bar.addSpacing(16)
        # Detection toggle & screenshot
        self.detection_toggle = QCheckBox("Enable Detection")
        self.detection_toggle.setChecked(True)
        self.detection_toggle.setStyleSheet("color: #fff; font-size: 14px;")
        self.detection_toggle.setEnabled(False)
        self.detection_toggle.toggled.connect(self.detection_toggled.emit)
        control_bar.addWidget(self.detection_toggle)
        self.screenshot_btn = QPushButton()
        self.screenshot_btn.setIcon(QIcon.fromTheme("camera-photo"))
        self.screenshot_btn.setText("Screenshot")
        self.screenshot_btn.setToolTip("Save current frame as image")
        self.screenshot_btn.setEnabled(False)
        self.screenshot_btn.setStyleSheet(self._button_style())
        self.screenshot_btn.clicked.connect(self.screenshot_clicked.emit)
        control_bar.addWidget(self.screenshot_btn)
        control_bar.addStretch(1)
        # Layout grid
        grid.addLayout(file_bar, 0, 0, 1, 1)
        grid.addWidget(video_frame, 1, 0, 1, 1)
        grid.addLayout(self.badge_bar, 2, 0, 1, 1)
        grid.addLayout(control_bar, 3, 0, 1, 1)
        grid.setRowStretch(1, 1)
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)

    def _button_style(self):
        return """
            QPushButton {
                background: #232323;
                border-radius: 24px;
                color: #fff;
                font-size: 15px;
                border: none;
            }
            QPushButton:hover {
                background: #03DAC5;
                color: #222;
            }
            QPushButton:pressed {
                background: #018786;
            }
        """

    def _select_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select Video File", "", "Video Files (*.mp4 *.avi *.mov *.mkv *.webm);;All Files (*)")
        if file_path:
            self.file_label.setText(file_path)
            self.file_selected.emit(file_path)
            self.video_loaded = True
            self._enable_controls(True)
            self.video_label.setText("")
            self.auto_select_model_device.emit()  # Request auto model/device selection

    def _enable_controls(self, enabled):
        self.play_btn.setEnabled(enabled)
        self.pause_btn.setEnabled(enabled)
        self.stop_btn.setEnabled(enabled)
        self.progress.setEnabled(enabled)
        self.detection_toggle.setEnabled(enabled)
        self.screenshot_btn.setEnabled(enabled)
        if enabled:
            self.auto_select_model_device.emit()  # Also trigger auto-select when controls are enabled

    def update_display(self, pixmap):
        # Maintain aspect ratio
        if pixmap:
            scaled = pixmap.scaled(self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
            self.video_label.setPixmap(scaled)
            self._set_controls_enabled(True)
            self.video_label.setStyleSheet("background: transparent; color: #888; font-size: 18px;")
        else:
            self.video_label.clear()
            self.video_label.setText("No video loaded. Please select a video file.")
            self._set_controls_enabled(False)
            self.video_label.setStyleSheet("background: transparent; color: #F44336; font-size: 18px;")

    def _set_controls_enabled(self, enabled):
        for btn in [self.play_btn, self.pause_btn, self.stop_btn, self.progress, self.detection_toggle, self.screenshot_btn]:
            btn.setEnabled(enabled)

    def update_stats(self, stats):
        # Accepts a stats dict for extensibility
        cars = stats.get('cars', 0)
        trucks = stats.get('trucks', 0)
        peds = stats.get('peds', 0)
        tlights = stats.get('tlights', 0)
        motorcycles = stats.get('motorcycles', 0)
        fps = stats.get('fps', None)
        # Try all possible keys for inference time
        inference = stats.get('inference', stats.get('detection_time', stats.get('detection_time_ms', None)))
        model = stats.get('model', stats.get('model_name', '-'))
        device = stats.get('device', stats.get('device_name', '-'))
        # Update overlay
        self.overlay.update_overlay(model, device, cars, trucks, peds, tlights, motorcycles)
        # Update FPS and Inference badges
        if fps is not None:
            self.fps_badge.setText(f"FPS: {fps:.2f}")
        else:
            self.fps_badge.setText("FPS: --")
        if inference is not None:
            self.inference_badge.setText(f"Inference: {inference:.1f} ms")
        else:
            self.inference_badge.setText("Inference: -- ms")

    def update_progress(self, value, max_value, timestamp):
        self.progress.setMaximum(max_value)
        self.progress.setValue(value)
        # Format timestamp as string (e.g., "00:00 / 00:00" or just str)
        if isinstance(timestamp, float) or isinstance(timestamp, int):
            timestamp_str = f"{timestamp:.2f}"
        else:
            timestamp_str = str(timestamp)
        self.timestamp.setText(timestamp_str)
