object-segmentation/src/gui/tabs/results_tab.py

"""
Results tab for browsing stored detections and visualizing overlays.
"""

from pathlib import Path
from typing import Dict, List, Optional

from PySide6.QtWidgets import (
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QLabel,
    QGroupBox,
    QPushButton,
    QSplitter,
    QTableWidget,
    QTableWidgetItem,
    QHeaderView,
    QAbstractItemView,
    QMessageBox,
    QCheckBox,
)
from PySide6.QtCore import Qt

from src.database.db_manager import DatabaseManager
from src.utils.config_manager import ConfigManager
from src.utils.logger import get_logger
from src.utils.image import Image, ImageLoadError
from src.gui.widgets import AnnotationCanvasWidget


logger = get_logger(__name__)


class ResultsTab(QWidget):
    """Results tab showing detection history and preview overlays."""

    def __init__(
        self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None
    ):
        super().__init__(parent)
        self.db_manager = db_manager
        self.config_manager = config_manager

        self.detection_summary: List[Dict] = []
        self.current_selection: Optional[Dict] = None
        self.current_image: Optional[Image] = None
        self.current_detections: List[Dict] = []
        self._image_path_cache: Dict[str, str] = {}

        self._setup_ui()
        self.refresh()

    def _setup_ui(self):
        """Setup user interface."""
        layout = QVBoxLayout()

        # Splitter for list + preview
        splitter = QSplitter(Qt.Horizontal)

        # Left pane: detection list
        left_container = QWidget()
        left_layout = QVBoxLayout()
        left_layout.setContentsMargins(0, 0, 0, 0)

        controls_layout = QHBoxLayout()
        self.refresh_btn = QPushButton("Refresh")
        self.refresh_btn.clicked.connect(self.refresh)
        controls_layout.addWidget(self.refresh_btn)
        controls_layout.addStretch()
        left_layout.addLayout(controls_layout)

        self.results_table = QTableWidget(0, 5)
        self.results_table.setHorizontalHeaderLabels(
            ["Image", "Model", "Detections", "Classes", "Last Updated"]
        )
        self.results_table.horizontalHeader().setSectionResizeMode(
            0, QHeaderView.Stretch
        )
        self.results_table.horizontalHeader().setSectionResizeMode(
            1, QHeaderView.Stretch
        )
        self.results_table.horizontalHeader().setSectionResizeMode(
            2, QHeaderView.ResizeToContents
        )
        self.results_table.horizontalHeader().setSectionResizeMode(
            3, QHeaderView.Stretch
        )
        self.results_table.horizontalHeader().setSectionResizeMode(
            4, QHeaderView.ResizeToContents
        )
        self.results_table.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.results_table.setSelectionMode(QAbstractItemView.SingleSelection)
        self.results_table.setEditTriggers(QAbstractItemView.NoEditTriggers)
        self.results_table.itemSelectionChanged.connect(self._on_result_selected)

        left_layout.addWidget(self.results_table)
        left_container.setLayout(left_layout)

        # Right pane: preview canvas and controls
        right_container = QWidget()
        right_layout = QVBoxLayout()
        right_layout.setContentsMargins(0, 0, 0, 0)

        preview_group = QGroupBox("Detection Preview")
        preview_layout = QVBoxLayout()

        self.preview_canvas = AnnotationCanvasWidget()
        self.preview_canvas.set_polyline_enabled(False)
        self.preview_canvas.set_show_bboxes(True)
        preview_layout.addWidget(self.preview_canvas)

        toggles_layout = QHBoxLayout()
        self.show_masks_checkbox = QCheckBox("Show Masks")
        self.show_masks_checkbox.setChecked(False)
        self.show_masks_checkbox.stateChanged.connect(self._apply_detection_overlays)
        self.show_bboxes_checkbox = QCheckBox("Show Bounding Boxes")
        self.show_bboxes_checkbox.setChecked(True)
        self.show_bboxes_checkbox.stateChanged.connect(self._toggle_bboxes)
        self.show_confidence_checkbox = QCheckBox("Show Confidence")
        self.show_confidence_checkbox.setChecked(False)
        self.show_confidence_checkbox.stateChanged.connect(
            self._apply_detection_overlays
        )
        toggles_layout.addWidget(self.show_masks_checkbox)
        toggles_layout.addWidget(self.show_bboxes_checkbox)
        toggles_layout.addWidget(self.show_confidence_checkbox)
        toggles_layout.addStretch()
        preview_layout.addLayout(toggles_layout)

        self.summary_label = QLabel("Select a detection result to preview.")
        self.summary_label.setWordWrap(True)
        preview_layout.addWidget(self.summary_label)

        preview_group.setLayout(preview_layout)
        right_layout.addWidget(preview_group)
        right_container.setLayout(right_layout)

        splitter.addWidget(left_container)
        splitter.addWidget(right_container)
        splitter.setStretchFactor(0, 1)
        splitter.setStretchFactor(1, 2)

        layout.addWidget(splitter)
        self.setLayout(layout)

    def refresh(self):
        """Refresh the detection list and preview."""
        self._load_detection_summary()
        self._populate_results_table()
        self.current_selection = None
        self.current_image = None
        self.current_detections = []
        self.preview_canvas.clear()
        self.summary_label.setText("Select a detection result to preview.")

    def _load_detection_summary(self):
        """Load latest detection summaries grouped by image + model."""
        try:
            detections = self.db_manager.get_detections(limit=500)
            summary_map: Dict[tuple, Dict] = {}

            for det in detections:
                key = (det["image_id"], det["model_id"])
                metadata = det.get("metadata") or {}
                entry = summary_map.setdefault(
                    key,
                    {
                        "image_id": det["image_id"],
                        "model_id": det["model_id"],
                        "image_path": det.get("image_path"),
                        "image_filename": det.get("image_filename")
                        or det.get("image_path"),
                        "model_name": det.get("model_name", ""),
                        "model_version": det.get("model_version", ""),
                        "last_detected": det.get("detected_at"),
                        "count": 0,
                        "classes": set(),
                        "source_path": metadata.get("source_path"),
                        "repository_root": metadata.get("repository_root"),
                    },
                )

                entry["count"] += 1
                if det.get("detected_at") and (
                    not entry.get("last_detected")
                    or str(det.get("detected_at")) > str(entry.get("last_detected"))
                ):
                    entry["last_detected"] = det.get("detected_at")
                if det.get("class_name"):
                    entry["classes"].add(det["class_name"])
                if metadata.get("source_path") and not entry.get("source_path"):
                    entry["source_path"] = metadata.get("source_path")
                if metadata.get("repository_root") and not entry.get("repository_root"):
                    entry["repository_root"] = metadata.get("repository_root")

            self.detection_summary = sorted(
                summary_map.values(),
                key=lambda x: str(x.get("last_detected") or ""),
                reverse=True,
            )
        except Exception as e:
            logger.error(f"Failed to load detection summary: {e}")
            QMessageBox.critical(
                self,
                "Error",
                f"Failed to load detection results:\n{str(e)}",
            )
            self.detection_summary = []

    def _populate_results_table(self):
        """Populate the table widget with detection summaries."""
        self.results_table.setRowCount(len(self.detection_summary))

        for row, entry in enumerate(self.detection_summary):
            model_label = f"{entry['model_name']} {entry['model_version']}".strip()
            class_list = (
                ", ".join(sorted(entry["classes"])) if entry["classes"] else "-"
            )

            items = [
                QTableWidgetItem(entry.get("image_filename", "")),
                QTableWidgetItem(model_label),
                QTableWidgetItem(str(entry.get("count", 0))),
                QTableWidgetItem(class_list),
                QTableWidgetItem(str(entry.get("last_detected") or "")),
            ]

            for col, item in enumerate(items):
                item.setData(Qt.UserRole, row)
                self.results_table.setItem(row, col, item)

        self.results_table.clearSelection()

    def _on_result_selected(self):
        """Handle selection changes in the detection table."""
        selected_items = self.results_table.selectedItems()
        if not selected_items:
            return

        row = selected_items[0].data(Qt.UserRole)
        if row is None or row >= len(self.detection_summary):
            return

        entry = self.detection_summary[row]
        if (
            self.current_selection
            and self.current_selection.get("image_id") == entry["image_id"]
            and self.current_selection.get("model_id") == entry["model_id"]
        ):
            return

        self.current_selection = entry

        image_path = self._resolve_image_path(entry)
        if not image_path:
            QMessageBox.warning(
                self,
                "Image Not Found",
                "Unable to locate the image file for this detection.",
            )
            return

        try:
            self.current_image = Image(image_path)
            self.preview_canvas.load_image(self.current_image)
        except ImageLoadError as e:
            logger.error(f"Failed to load image '{image_path}': {e}")
            QMessageBox.critical(
                self,
                "Image Error",
                f"Failed to load image for preview:\n{str(e)}",
            )
            return

        self._load_detections_for_selection(entry)
        self._apply_detection_overlays()
        self._update_summary_label(entry)

    def _load_detections_for_selection(self, entry: Dict):
        """Load detection records for the selected image/model pair."""
        self.current_detections = []
        if not entry:
            return

        try:
            filters = {"image_id": entry["image_id"], "model_id": entry["model_id"]}
            self.current_detections = self.db_manager.get_detections(filters)
        except Exception as e:
            logger.error(f"Failed to load detections for preview: {e}")
            QMessageBox.critical(
                self,
                "Error",
                f"Failed to load detections for this image:\n{str(e)}",
            )
            self.current_detections = []

    def _apply_detection_overlays(self):
        """Draw detections onto the preview canvas based on current toggles."""
        self.preview_canvas.clear_annotations()
        self.preview_canvas.set_show_bboxes(self.show_bboxes_checkbox.isChecked())

        if not self.current_detections or not self.current_image:
            return

        for det in self.current_detections:
            color = self._get_class_color(det.get("class_name"))

            if self.show_masks_checkbox.isChecked() and det.get("segmentation_mask"):
                mask_points = self._convert_mask(det["segmentation_mask"])
                if mask_points:
                    self.preview_canvas.draw_saved_polyline(mask_points, color)

            bbox = [
                det.get("x_min"),
                det.get("y_min"),
                det.get("x_max"),
                det.get("y_max"),
            ]
            if all(v is not None for v in bbox):
                label = None
                if self.show_confidence_checkbox.isChecked():
                    confidence = det.get("confidence")
                    if confidence is not None:
                        label = f"{confidence:.2f}"
                self.preview_canvas.draw_saved_bbox(bbox, color, label=label)

    def _convert_mask(self, mask_points: List[List[float]]) -> List[List[float]]:
        """Convert stored [x, y] masks to [y, x] format for the canvas."""
        converted = []
        for point in mask_points:
            if len(point) >= 2:
                x, y = point[0], point[1]
                converted.append([y, x])
        return converted

    def _toggle_bboxes(self):
        """Update bounding box visibility on the canvas."""
        self.preview_canvas.set_show_bboxes(self.show_bboxes_checkbox.isChecked())
        # Re-render to respect show/hide when toggled
        self._apply_detection_overlays()

    def _update_summary_label(self, entry: Dict):
        """Display textual summary for the selected detection run."""
        classes = ", ".join(sorted(entry.get("classes", []))) or "-"
        summary_text = (
            f"Image: {entry.get('image_filename', 'unknown')}\n"
            f"Model: {entry.get('model_name', '')} {entry.get('model_version', '')}\n"
            f"Detections: {entry.get('count', 0)}\n"
            f"Classes: {classes}\n"
            f"Last Updated: {entry.get('last_detected', 'n/a')}"
        )
        self.summary_label.setText(summary_text)

    def _resolve_image_path(self, entry: Dict) -> Optional[str]:
        """Resolve an image path using metadata, cache, and repository hints."""
        relative_path = entry.get("image_path") if entry else None
        cache_key = relative_path or entry.get("source_path")
        if cache_key and cache_key in self._image_path_cache:
            cached = Path(self._image_path_cache[cache_key])
            if cached.exists():
                return self._image_path_cache[cache_key]
            del self._image_path_cache[cache_key]

        candidates = []
        source_path = entry.get("source_path") if entry else None
        if source_path:
            candidates.append(Path(source_path))

        repo_roots = []
        if entry.get("repository_root"):
            repo_roots.append(entry["repository_root"])
        config_repo = self.config_manager.get_image_repository_path()
        if config_repo:
            repo_roots.append(config_repo)

        for root in repo_roots:
            if relative_path:
                candidates.append(Path(root) / relative_path)

        if relative_path:
            candidates.append(Path(relative_path))

        for candidate in candidates:
            try:
                if candidate and candidate.exists():
                    resolved = str(candidate.resolve())
                    if cache_key:
                        self._image_path_cache[cache_key] = resolved
                    return resolved
            except Exception:
                continue

        # Fallback: search by filename in known roots
        filename = Path(relative_path).name if relative_path else None
        if filename:
            search_roots = [Path(root) for root in repo_roots if root]
            if not search_roots:
                search_roots = [Path("data")]
            match = self._search_in_roots(filename, search_roots)
            if match:
                resolved = str(match.resolve())
                if cache_key:
                    self._image_path_cache[cache_key] = resolved
                return resolved

        return None

    def _search_in_roots(self, filename: str, roots: List[Path]) -> Optional[Path]:
        """Search for a file name within a list of root directories."""
        for root in roots:
            try:
                if not root.exists():
                    continue
                for candidate in root.rglob(filename):
                    return candidate
            except Exception as e:
                logger.debug(f"Error searching for {filename} in {root}: {e}")
        return None

    def _get_class_color(self, class_name: Optional[str]) -> str:
        """Return consistent color hex for a class name."""
        if not class_name:
            return "#FF6B6B"

        color_map = self.config_manager.get_bbox_colors()
        if class_name in color_map:
            return color_map[class_name]

        # Deterministic fallback color based on hash
        palette = [
            "#FF6B6B",
            "#4ECDC4",
            "#FFD166",
            "#1D3557",
            "#F4A261",
            "#E76F51",
        ]
        return palette[hash(class_name) % len(palette)]