app/region_selector.py

"""
RegionSelectorDialog — Dialog chọn vùng che sub với scrubber toàn video.

Nâng cấp so với bản cũ:
- Thanh scrubber QSlider trượt toàn bộ video để kiểm tra vùng che ở mọi frame
- Nút Play/Pause để xem tự động (25fps qua QTimer + seek, không decode realtime)
- Phím tắt: Space=Play/Dừng, ←=Lùi frame, →=Tiến frame, Enter=Lưu
- Overlay: vùng đã lưu fill đỏ mờ + label "✓ Vùng che"; đang vẽ fill vàng nét đứt
- Giữ VideoCapture mở xuyên suốt dialog → seek nhanh; release khi đóng

Backward compat: API giống bản cũ (constructor, self.region, accept/reject).
"""

import cv2
from PyQt6.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QSlider, QMessageBox
)
from PyQt6.QtCore import Qt, QTimer
from PyQt6.QtGui import QImage, QPixmap, QPainter, QPen, QColor


class RegionSelectorDialog(QDialog):
    """
    Dialog chọn vùng che sub.

    Kết quả:
        self.region = (x, y, w, h, orig_video_w, orig_video_h)  — pixel gốc video
        self.region = None nếu chưa vẽ hoặc bỏ qua
    """

    def __init__(
        self,
        video_path,
        parent=None,
        is_batch: bool = False,
        current_idx: int = 1,
        total_videos: int = 1,
    ):
        super().__init__(parent)

        if is_batch:
            self.setWindowTitle(
                f"[{current_idx}/{total_videos}] Kéo thả chọn vùng che Sub "
                "— Trượt / Phát để kiểm tra toàn video"
            )
        else:
            self.setWindowTitle(
                "Kéo thả chuột để chọn vùng che Sub "
                "— Trượt / Phát để kiểm tra toàn video (X để lưu)"
            )

        self.setModal(True)
        self.video_path = str(video_path)
        self.region = None

        # ── Mở VideoCapture (giữ mở suốt dialog) ────────────────────────
        self._cap = cv2.VideoCapture(self.video_path)
        if not self._cap.isOpened():
            QMessageBox.critical(self, "Lỗi", f"Không mở được video:\n{self.video_path}")
            super().reject()
            return

        self._total_frames = max(1, int(self._cap.get(cv2.CAP_PROP_FRAME_COUNT)))
        self._fps = self._cap.get(cv2.CAP_PROP_FPS) or 25.0
        self.orig_w = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.orig_h = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # Load frame giữa video
        mid = self._total_frames // 2
        frame = self._read_frame(mid)
        if frame is None:
            QMessageBox.critical(self, "Lỗi", "Không thể đọc khung hình video!")
            self._cap.release()
            super().reject()
            return

        # ── Display state ────────────────────────────────────────────────
        self.MAX_W = 1280
        self.MAX_H = 720
        self._current_frame_idx = mid
        self.disp_w = 1
        self.disp_h = 1
        self.scale_x = 1.0
        self.scale_y = 1.0

        # ── Drawing state ────────────────────────────────────────────────
        self.start_point = None
        self.end_point = None
        self.drawing = False

        # ── Widgets ──────────────────────────────────────────────────────
        self.label = QLabel()
        self.label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.label.setCursor(Qt.CursorShape.CrossCursor)
        self.label.mousePressEvent = self.mouse_press
        self.label.mouseMoveEvent = self.mouse_move
        self.label.mouseReleaseEvent = self.mouse_release
        self.label.paintEvent = self.paint_label

        hint = QLabel(
            "📢 Kéo chuột vẽ vùng che (hình chữ nhật đỏ).  "
            "▸ Trượt thanh để kiểm tra mọi frame.  "
            "▶ Nút Phát xem tự động.  "
            "[Space]=Play/Dừng  [←→]=Lùi/Tiến frame  [Enter]=Lưu"
        )
        hint.setWordWrap(True)
        hint.setStyleSheet("font-size:12px; padding:4px; color:#ccc; background:#1a1a1e;")

        # ── Scrubber row ─────────────────────────────────────────────────
        self._playing = False

        self.btn_play = QPushButton("▶ Phát")
        self.btn_play.setFixedWidth(100)
        self._style_play_btn(playing=False)
        self.btn_play.clicked.connect(self._toggle_play)

        self.lbl_time = QLabel("00:00:00")
        self.lbl_time.setFixedWidth(72)
        self.lbl_time.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.lbl_time.setStyleSheet("font-family:monospace; font-size:12px; color:#aaa;")

        self.slider = QSlider(Qt.Orientation.Horizontal)
        self.slider.setMinimum(0)
        self.slider.setMaximum(self._total_frames - 1)
        self.slider.setValue(mid)
        self.slider.setStyleSheet("""
            QSlider::groove:horizontal {
                height: 6px; background: #3a3a3a; border-radius: 3px;
            }
            QSlider::handle:horizontal {
                background: #e74c3c; border: none;
                width: 14px; height: 14px; margin: -4px 0; border-radius: 7px;
            }
            QSlider::sub-page:horizontal {
                background: #c0392b; border-radius: 3px;
            }
        """)
        self.slider.valueChanged.connect(self._on_slider_changed)
        self.slider.sliderPressed.connect(self._pause)

        total_sec = self._total_frames / self._fps
        mm, ss = divmod(int(total_sec), 60)
        hh, mm = divmod(mm, 60)
        lbl_total = QLabel(f"{hh:02d}:{mm:02d}:{ss:02d}")
        lbl_total.setAlignment(Qt.AlignmentFlag.AlignCenter)
        lbl_total.setStyleSheet("font-family:monospace; font-size:12px; color:#555;")
        lbl_total.setFixedWidth(72)

        scrub_row = QHBoxLayout()
        scrub_row.addWidget(self.btn_play)
        scrub_row.addWidget(self.lbl_time)
        scrub_row.addWidget(self.slider, stretch=1)
        scrub_row.addWidget(lbl_total)

        # ── Play timer ────────────────────────────────────────────────────
        self._timer = QTimer(self)
        self._timer.setInterval(int(1000 / min(self._fps, 25)))
        self._timer.timeout.connect(self._advance_frame)

        # ── Main layout ───────────────────────────────────────────────────
        layout = QVBoxLayout(self)
        layout.setSpacing(6)
        layout.addWidget(hint)
        layout.addWidget(self.label)
        layout.addLayout(scrub_row)

        # ── Batch mode buttons ────────────────────────────────────────────
        self.is_batch = is_batch
        if is_batch:
            btn_layout = QHBoxLayout()

            self.btn_skip = QPushButton("⏭️ Bỏ qua (Không che)")
            self.btn_skip.setStyleSheet(
                "background-color:#7F8C8D; color:white; padding:10px; "
                "font-weight:bold; font-size:14px;"
            )
            self.btn_skip.clicked.connect(self._reject_and_stop)

            label_next = "✅ OK, Video cuối!" if current_idx == total_videos else "✅ OK, đã che xong → Next"
            self.btn_next = QPushButton(label_next)
            self.btn_next.setStyleSheet(
                "background-color:#27AE60; color:white; padding:10px; "
                "font-weight:bold; font-size:14px;"
            )
            self.btn_next.clicked.connect(self._accept_and_stop)

            btn_layout.addWidget(self.btn_skip)
            btn_layout.addWidget(self.btn_next)
            layout.addLayout(btn_layout)

        # Load frame đầu tiên
        self._set_frame(mid)

    # ── Video helpers ─────────────────────────────────────────────────────

    def _read_frame(self, frame_idx: int):
        """Seek và đọc frame, trả về numpy RGB hoặc None."""
        self._cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        ret, frame = self._cap.read()
        if not ret or frame is None:
            return None
        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    def _set_frame(self, frame_idx: int):
        """Seek đến frame_idx, cập nhật label + slider + timestamp."""
        frame_idx = max(0, min(frame_idx, self._total_frames - 1))
        self._current_frame_idx = frame_idx

        frame = self._read_frame(frame_idx)
        if frame is None:
            return

        h, w, ch = frame.shape
        scaled = QPixmap.fromImage(
            QImage(frame.data, w, h, ch * w, QImage.Format.Format_RGB888)
        ).scaled(
            self.MAX_W, self.MAX_H,
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation,
        )
        self.disp_w = scaled.width()
        self.disp_h = scaled.height()
        self.scale_x = self.orig_w / self.disp_w
        self.scale_y = self.orig_h / self.disp_h

        self.label.setPixmap(scaled)
        self.label.setFixedSize(self.disp_w, self.disp_h)
        self.label.update()

        self.slider.blockSignals(True)
        self.slider.setValue(frame_idx)
        self.slider.blockSignals(False)

        sec = frame_idx / self._fps
        mm, ss = divmod(int(sec), 60)
        hh, mm = divmod(mm, 60)
        self.lbl_time.setText(f"{hh:02d}:{mm:02d}:{ss:02d}")

    # ── Playback ──────────────────────────────────────────────────────────

    def _style_play_btn(self, playing: bool):
        if playing:
            self.btn_play.setText("⏸ Tạm dừng")
            self.btn_play.setStyleSheet(
                "background:#e67e22; color:#000; font-weight:bold; border-radius:4px; padding:4px;"
            )
        else:
            self.btn_play.setText("▶ Phát")
            self.btn_play.setStyleSheet(
                "background:#2ecc71; color:#000; font-weight:bold; border-radius:4px; padding:4px;"
            )

    def _on_slider_changed(self, value: int):
        self._set_frame(value)

    def _toggle_play(self):
        if self._playing:
            self._pause()
        else:
            self._play()

    def _play(self):
        self._playing = True
        self._style_play_btn(playing=True)
        self._timer.start()

    def _pause(self):
        self._playing = False
        self._timer.stop()
        self._style_play_btn(playing=False)

    def _advance_frame(self):
        """Timer tick: tiến 1 frame, loop khi hết."""
        nxt = self._current_frame_idx + 1
        if nxt >= self._total_frames:
            nxt = 0
        self._set_frame(nxt)

    # ── Coordinate normalization ──────────────────────────────────────────

    def _display_geometry(self):
        """Return (off_x, off_y, disp_w, disp_h) — where the video pixmap actually
        sits inside the QLabel, accounting for any centered letterbox/pillarbox
        margin so mouse coordinates are never measured against the black bars."""
        label_w = self.label.width()
        label_h = self.label.height()
        off_x = max(0, (label_w - self.disp_w) // 2)
        off_y = max(0, (label_h - self.disp_h) // 2)
        return off_x, off_y, self.disp_w, self.disp_h

    def _ui_to_native(self, ux, uy):
        """Convert UI (label-local) coords -> normalized (0.0..1.0) coords ->
        native video frame coords. Logs the full chain for debugging:
        [UI Cords] -> [Normalized Cords] -> [Native Frame Cords]."""
        off_x, off_y, dw, dh = self._display_geometry()
        px = max(0.0, min(float(ux - off_x), float(dw)))
        py = max(0.0, min(float(uy - off_y), float(dh)))
        norm_x = px / dw if dw > 0 else 0.0
        norm_y = py / dh if dh > 0 else 0.0
        real_x = int(round(norm_x * self.orig_w))
        real_y = int(round(norm_y * self.orig_h))
        print(
            f"[UI Cords] ({ux:.1f}, {uy:.1f}) -> "
            f"[Normalized Cords] ({norm_x:.4f}, {norm_y:.4f}) -> "
            f"[Native Frame Cords] ({real_x}, {real_y})"
        )
        return real_x, real_y

    def _native_to_display(self, nx, ny):
        """Convert native frame coords back to label-local display coords."""
        off_x, off_y, dw, dh = self._display_geometry()
        dx = off_x + int(round(nx / self.orig_w * dw)) if self.orig_w else off_x
        dy = off_y + int(round(ny / self.orig_h * dh)) if self.orig_h else off_y
        return dx, dy

    # ── Drawing ───────────────────────────────────────────────────────────

    def mouse_press(self, event):
        self._pause()
        self.start_point = event.position().toPoint()
        self.end_point = self.start_point
        self.drawing = True
        self.label.update()

    def mouse_move(self, event):
        if self.drawing:
            self.end_point = event.position().toPoint()
            self.label.update()

    def mouse_release(self, event):
        self.end_point = event.position().toPoint()
        self.drawing = False
        self.label.update()

        off_x, off_y, dw, dh = self._display_geometry()
        x1 = max(off_x, min(min(self.start_point.x(), self.end_point.x()), off_x + dw))
        y1 = max(off_y, min(min(self.start_point.y(), self.end_point.y()), off_y + dh))
        x2 = max(off_x, min(max(self.start_point.x(), self.end_point.x()), off_x + dw))
        y2 = max(off_y, min(max(self.start_point.y(), self.end_point.y()), off_y + dh))

        orig_x, orig_y = self._ui_to_native(x1, y1)
        rx2, ry2 = self._ui_to_native(x2, y2)
        orig_rw = max(1, rx2 - orig_x)
        orig_rh = max(1, ry2 - orig_y)

        if orig_rw > 0 and orig_rh > 0:
            self.region = (orig_x, orig_y, orig_rw, orig_rh, self.orig_w, self.orig_h)
            print(
                f"[REGION] stored native region: x={orig_x}, y={orig_y}, "
                f"w={orig_rw}, h={orig_rh} (frame {self.orig_w}x{self.orig_h})"
            )

    def keyPressEvent(self, event):
        k = event.key()
        if k in (Qt.Key.Key_Return, Qt.Key.Key_Enter):
            self._accept_and_stop()
        elif k == Qt.Key.Key_Space:
            self._toggle_play()
        elif k == Qt.Key.Key_Left:
            self._pause()
            self._set_frame(self._current_frame_idx - 1)
        elif k == Qt.Key.Key_Right:
            self._pause()
            self._set_frame(self._current_frame_idx + 1)
        else:
            super().keyPressEvent(event)

    def paint_label(self, event):
        """Overlay: vùng đã lưu (đỏ mờ) + đang vẽ (vàng nét đứt)."""
        QLabel.paintEvent(self.label, event)
        painter = QPainter(self.label)
        _off_x, _off_y, _dw, _dh = self._display_geometry()

        # Vùng đã lưu
        if self.region:
            rx, ry, rw, rh, _vw, _vh = self.region
            dx, dy = self._native_to_display(rx, ry)
            dx2, dy2 = self._native_to_display(rx + rw, ry + rh)
            dw = dx2 - dx
            dh = dy2 - dy
            painter.fillRect(dx, dy, dw, dh, QColor(220, 0, 0, 70))
            painter.setPen(QPen(QColor(255, 60, 60), 2, Qt.PenStyle.SolidLine))
            painter.drawRect(dx, dy, dw, dh)
            painter.setPen(QColor(255, 220, 0))
            painter.drawText(dx + 4, dy + 15, "✓ Vùng che")

        # Đang vẽ
        if self.start_point and self.end_point and self.drawing:
            x1 = min(self.start_point.x(), self.end_point.x())
            y1 = min(self.start_point.y(), self.end_point.y())
            w = abs(self.start_point.x() - self.end_point.x())
            h = abs(self.start_point.y() - self.end_point.y())
            painter.fillRect(x1, y1, w, h, QColor(255, 255, 0, 45))
            painter.setPen(QPen(QColor(255, 230, 0), 1, Qt.PenStyle.DashLine))
            painter.drawRect(x1, y1, w, h)

        painter.end()

    # ── Lifecycle ─────────────────────────────────────────────────────────

    def _release(self):
        self._pause()
        if hasattr(self, "_cap") and self._cap.isOpened():
            self._cap.release()

    def _accept_and_stop(self):
        self._release()
        super().accept()

    def _reject_and_stop(self):
        self._release()
        super().reject()

    def closeEvent(self, event):
        self._release()
        super().closeEvent(event)

    def accept(self):
        self._release()
        super().accept()

    def reject(self):
        self._release()
        super().reject()