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, Tuple

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)

        self.delete_all_btn = QPushButton("Delete All Detections")
        self.delete_all_btn.setToolTip(
            "Permanently delete ALL detections from the database.\n" "This cannot be undone."
        )
        self.delete_all_btn.clicked.connect(self._delete_all_detections)
        controls_layout.addWidget(self.delete_all_btn)

        self.export_labels_btn = QPushButton("Export Labels")
        self.export_labels_btn.setToolTip(
            "Export YOLO .txt labels for the selected image/model run.\n"
            "Output path is inferred from the image path (images/ -> labels/)."
        )
        self.export_labels_btn.clicked.connect(self._export_labels_for_current_selection)
        controls_layout.addWidget(self.export_labels_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()
        # Auto-zoom so newly loaded images fill the available preview viewport.
        self.preview_canvas.set_auto_fit_to_view(True)
        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 _delete_all_detections(self):
        """Delete all detections from the database after user confirmation."""
        confirm = QMessageBox.warning(
            self,
            "Delete All Detections",
            "This will permanently delete ALL detections from the database.\n\n"
            "This action cannot be undone.\n\n"
            "Do you want to continue?",
            QMessageBox.Yes | QMessageBox.No,
            QMessageBox.No,
        )

        if confirm != QMessageBox.Yes:
            return

        try:
            deleted = self.db_manager.delete_all_detections()
        except Exception as exc:
            logger.error(f"Failed to delete all detections: {exc}")
            QMessageBox.critical(
                self,
                "Error",
                f"Failed to delete detections:\n{exc}",
            )
            return

        QMessageBox.information(
            self,
            "Delete All Detections",
            f"Deleted {deleted} detection(s) from the database.",
        )

        # Reset UI state.
        self.refresh()

    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.")
        if hasattr(self, "export_labels_btn"):
            self.export_labels_btn.setEnabled(False)

    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)
        if hasattr(self, "export_labels_btn"):
            self.export_labels_btn.setEnabled(True)

    def _export_labels_for_current_selection(self):
        """Export YOLO label file(s) for the currently selected image/model."""
        if not self.current_selection:
            QMessageBox.information(self, "Export Labels", "Select a detection result first.")
            return

        entry = self.current_selection

        image_path_str = self._resolve_image_path(entry)
        if not image_path_str:
            QMessageBox.warning(
                self,
                "Export Labels",
                "Unable to locate the image file for this detection; cannot infer labels path.",
            )
            return

        # Ensure we have the detections for the selection.
        if not self.current_detections:
            self._load_detections_for_selection(entry)

        if not self.current_detections:
            QMessageBox.information(
                self,
                "Export Labels",
                "No detections found for this image/model selection.",
            )
            return

        image_path = Path(image_path_str)
        try:
            label_path = self._infer_yolo_label_path(image_path)
        except Exception as exc:
            logger.error(f"Failed to infer label path for {image_path}: {exc}")
            QMessageBox.critical(
                self,
                "Export Labels",
                f"Failed to infer export path for labels:\n{exc}",
            )
            return

        class_map = self._build_detection_class_index_map(self.current_detections)
        if not class_map:
            QMessageBox.warning(
                self,
                "Export Labels",
                "Unable to build class->index mapping (missing class names).",
            )
            return

        lines_written = 0
        skipped = 0
        label_path.parent.mkdir(parents=True, exist_ok=True)
        try:
            with open(label_path, "w", encoding="utf-8") as handle:
                print("writing to", label_path)
                for det in self.current_detections:
                    yolo_line = self._format_detection_as_yolo_line(det, class_map)
                    if not yolo_line:
                        skipped += 1
                        continue
                    handle.write(yolo_line + "\n")
                    lines_written += 1
        except OSError as exc:
            logger.error(f"Failed to write labels file {label_path}: {exc}")
            QMessageBox.critical(
                self,
                "Export Labels",
                f"Failed to write label file:\n{label_path}\n\n{exc}",
            )
            return

        return
        # Optional: write a classes.txt next to the labels root to make the mapping discoverable.
        # This is not required by Ultralytics (data.yaml usually holds class names), but helps reuse.
        try:
            classes_txt = label_path.parent.parent / "classes.txt"
            classes_txt.parent.mkdir(parents=True, exist_ok=True)
            inv = {idx: name for name, idx in class_map.items()}
            with open(classes_txt, "w", encoding="utf-8") as handle:
                for idx in range(len(inv)):
                    handle.write(f"{inv[idx]}\n")
        except Exception:
            # Non-fatal
            pass

        QMessageBox.information(
            self,
            "Export Labels",
            f"Exported {lines_written} label line(s) to:\n{label_path}\n\nSkipped {skipped} invalid detection(s).",
        )

    def _infer_yolo_label_path(self, image_path: Path) -> Path:
        """Infer a YOLO label path from an image path.

        If the image lives under an `images/` directory (anywhere in the path), we mirror the
        subpath under a sibling `labels/` directory at the same level.

        Example:
            /dataset/train/images/sub/img.jpg -> /dataset/train/labels/sub/img.txt
        """

        resolved = image_path.expanduser().resolve()

        # Find the nearest ancestor directory named 'images'
        images_dir: Optional[Path] = None
        for parent in [resolved.parent, *resolved.parents]:
            if parent.name.lower() == "images":
                images_dir = parent
                break

        if images_dir is not None:
            rel = resolved.relative_to(images_dir)
            labels_dir = images_dir.parent / "labels"
            return (labels_dir / rel).with_suffix(".txt")

        # Fallback: create a local sibling labels folder next to the image.
        return (resolved.parent / "labels" / resolved.name).with_suffix(".txt")

    def _build_detection_class_index_map(self, detections: List[Dict]) -> Dict[str, int]:
        """Build a stable class_name -> YOLO class index mapping.

        Preference order:
        1) Database object_classes table (alphabetical class_name order)
        2) Fallback to class_name values present in the detections (alphabetical)
        """

        names: List[str] = []
        try:
            db_classes = self.db_manager.get_object_classes() or []
            names = [str(row.get("class_name")) for row in db_classes if row.get("class_name")]
        except Exception:
            names = []

        if not names:
            observed = sorted({str(det.get("class_name")) for det in detections if det.get("class_name")})
            names = list(observed)

        return {name: idx for idx, name in enumerate(names)}

    def _format_detection_as_yolo_line(self, det: Dict, class_map: Dict[str, int]) -> Optional[str]:
        """Convert a detection row to a YOLO label line.

        - If segmentation_mask is present, exports segmentation polygon format:
            class x1 y1 x2 y2 ...
          (normalized coordinates)
        - Otherwise exports bbox format:
            class x_center y_center width height
          (normalized coordinates)
        """

        class_name = det.get("class_name")
        if not class_name or class_name not in class_map:
            return None
        class_idx = class_map[class_name]

        mask = det.get("segmentation_mask")
        polygon = self._convert_segmentation_mask_to_polygon(mask)
        if polygon:
            coords = " ".join(f"{value:.6f}" for value in polygon)
            return f"{class_idx} {coords}".strip()

        bbox = self._convert_bbox_to_yolo_xywh(det)
        if bbox is None:
            return None
        x_center, y_center, width, height = bbox
        return f"{class_idx} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"

    def _convert_bbox_to_yolo_xywh(self, det: Dict) -> Optional[Tuple[float, float, float, float]]:
        """Convert stored xyxy (normalized) bbox to YOLO xywh (normalized)."""

        x_min = det.get("x_min")
        y_min = det.get("y_min")
        x_max = det.get("x_max")
        y_max = det.get("y_max")
        if any(v is None for v in (x_min, y_min, x_max, y_max)):
            return None

        try:
            x_min_f = self._clamp01(float(x_min))
            y_min_f = self._clamp01(float(y_min))
            x_max_f = self._clamp01(float(x_max))
            y_max_f = self._clamp01(float(y_max))
        except (TypeError, ValueError):
            return None

        width = max(0.0, x_max_f - x_min_f)
        height = max(0.0, y_max_f - y_min_f)
        if width <= 0.0 or height <= 0.0:
            return None

        x_center = x_min_f + width / 2.0
        y_center = y_min_f + height / 2.0
        return x_center, y_center, width, height

    def _convert_segmentation_mask_to_polygon(self, mask_data) -> List[float]:
        """Convert stored segmentation_mask [[x,y], ...] to YOLO polygon coords [x1,y1,...]."""

        if not isinstance(mask_data, list):
            return []

        coords: List[float] = []
        for point in mask_data:
            if not isinstance(point, (list, tuple)) or len(point) < 2:
                continue
            try:
                x = self._clamp01(float(point[0]))
                y = self._clamp01(float(point[1]))
            except (TypeError, ValueError):
                continue
            coords.extend([x, y])

        # Need at least 3 points => 6 values.
        return coords if len(coords) >= 6 else []

    @staticmethod
    def _clamp01(value: float) -> float:
        if value < 0.0:
            return 0.0
        if value > 1.0:
            return 1.0
        return value

    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)]