object-segmentation/src/gui/widgets/annotation_tools_widget.py

"""
Annotation tools widget for controlling annotation parameters.
Includes polyline tool, color picker, class selection, and annotation management.
"""

from PySide6.QtWidgets import (
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QLabel,
    QGroupBox,
    QPushButton,
    QComboBox,
    QSpinBox,
    QDoubleSpinBox,
    QCheckBox,
    QColorDialog,
    QInputDialog,
    QMessageBox,
)
from PySide6.QtGui import QColor, QIcon, QPixmap, QPainter
from PySide6.QtCore import Qt, Signal
from typing import Optional, Dict

from src.database.db_manager import DatabaseManager
from src.utils.logger import get_logger

logger = get_logger(__name__)


class AnnotationToolsWidget(QWidget):
    """
    Widget for annotation tool controls.

    Features:
    - Enable/disable polyline tool
    - Color selection for polyline pen
    - Object class selection
    - Add new object classes
    - Pen width control
    - Clear annotations

    Signals:
        polyline_enabled_changed: Emitted when polyline tool is enabled/disabled (bool)
        polyline_pen_color_changed: Emitted when polyline pen color changes (QColor)
        polyline_pen_width_changed: Emitted when polyline pen width changes (int)
        class_selected: Emitted when object class is selected (dict)
        clear_annotations_requested: Emitted when clear button is pressed
    """

    polyline_enabled_changed = Signal(bool)
    polyline_pen_color_changed = Signal(QColor)
    polyline_pen_width_changed = Signal(int)
    simplify_on_finish_changed = Signal(bool)
    simplify_epsilon_changed = Signal(float)
    # Toggle visibility of bounding boxes on the canvas
    show_bboxes_changed = Signal(bool)
    class_selected = Signal(dict)
    class_color_changed = Signal()
    clear_annotations_requested = Signal()

    def __init__(self, db_manager: DatabaseManager, parent=None):
        """
        Initialize annotation tools widget.

        Args:
            db_manager: Database manager instance
            parent: Parent widget
        """
        super().__init__(parent)
        self.db_manager = db_manager
        self.polyline_enabled = False
        self.current_color = QColor(255, 0, 0, 128)  # Red with 50% alpha
        self.current_class = None

        self._setup_ui()
        self._load_object_classes()

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

        # Polyline Tool Group
        polyline_group = QGroupBox("Polyline Tool")
        polyline_layout = QVBoxLayout()

        # Enable/Disable polyline tool
        button_layout = QHBoxLayout()
        self.polyline_toggle_btn = QPushButton("Start Drawing Polyline")
        self.polyline_toggle_btn.setCheckable(True)
        self.polyline_toggle_btn.clicked.connect(self._on_polyline_toggle)
        button_layout.addWidget(self.polyline_toggle_btn)
        polyline_layout.addLayout(button_layout)

        # Polyline pen width control
        width_layout = QHBoxLayout()
        width_layout.addWidget(QLabel("Pen Width:"))
        self.polyline_pen_width_spin = QSpinBox()
        self.polyline_pen_width_spin.setMinimum(1)
        self.polyline_pen_width_spin.setMaximum(20)
        self.polyline_pen_width_spin.setValue(3)
        self.polyline_pen_width_spin.valueChanged.connect(
            self._on_polyline_pen_width_changed
        )
        width_layout.addWidget(self.polyline_pen_width_spin)
        width_layout.addStretch()
        polyline_layout.addLayout(width_layout)

        # Simplification controls (RDP)
        simplify_layout = QHBoxLayout()
        self.simplify_checkbox = QCheckBox("Simplify on finish")
        self.simplify_checkbox.setChecked(True)
        self.simplify_checkbox.stateChanged.connect(self._on_simplify_toggle)
        simplify_layout.addWidget(self.simplify_checkbox)

        simplify_layout.addWidget(QLabel("epsilon (px):"))
        self.eps_spin = QDoubleSpinBox()
        self.eps_spin.setRange(0.0, 1000.0)
        self.eps_spin.setSingleStep(0.5)
        self.eps_spin.setValue(2.0)
        self.eps_spin.valueChanged.connect(self._on_eps_change)
        simplify_layout.addWidget(self.eps_spin)
        simplify_layout.addStretch()
        polyline_layout.addLayout(simplify_layout)

        polyline_group.setLayout(polyline_layout)
        layout.addWidget(polyline_group)

        # Object Class Group
        class_group = QGroupBox("Object Class")
        class_layout = QVBoxLayout()

        # Class selection dropdown
        self.class_combo = QComboBox()
        self.class_combo.currentIndexChanged.connect(self._on_class_selected)
        class_layout.addWidget(self.class_combo)

        # Add / manage classes
        class_button_layout = QHBoxLayout()
        self.add_class_btn = QPushButton("Add New Class")
        self.add_class_btn.clicked.connect(self._on_add_class)
        class_button_layout.addWidget(self.add_class_btn)

        self.refresh_classes_btn = QPushButton("Refresh")
        self.refresh_classes_btn.clicked.connect(self._load_object_classes)
        class_button_layout.addWidget(self.refresh_classes_btn)
        class_layout.addLayout(class_button_layout)

        # Class color (associated with selected object class)
        color_layout = QHBoxLayout()
        color_layout.addWidget(QLabel("Class Color:"))
        self.color_btn = QPushButton()
        self.color_btn.setFixedSize(40, 30)
        self.color_btn.clicked.connect(self._on_color_picker)
        self._update_color_button()
        color_layout.addWidget(self.color_btn)
        color_layout.addStretch()
        class_layout.addLayout(color_layout)

        # Selected class info
        self.class_info_label = QLabel("No class selected")
        self.class_info_label.setWordWrap(True)
        self.class_info_label.setStyleSheet(
            "QLabel { color: #888; font-style: italic; }"
        )
        class_layout.addWidget(self.class_info_label)

        class_group.setLayout(class_layout)
        layout.addWidget(class_group)

        # Actions Group
        actions_group = QGroupBox("Actions")
        actions_layout = QVBoxLayout()

        # Show / hide bounding boxes
        self.show_bboxes_checkbox = QCheckBox("Show bounding boxes")
        self.show_bboxes_checkbox.setChecked(True)
        self.show_bboxes_checkbox.stateChanged.connect(self._on_show_bboxes_toggle)
        actions_layout.addWidget(self.show_bboxes_checkbox)

        self.clear_btn = QPushButton("Clear All Annotations")
        self.clear_btn.clicked.connect(self._on_clear_annotations)
        actions_layout.addWidget(self.clear_btn)

        actions_group.setLayout(actions_layout)
        layout.addWidget(actions_group)

        layout.addStretch()
        self.setLayout(layout)

    def _update_color_button(self):
        """Update the color button appearance with current color."""
        pixmap = QPixmap(40, 30)
        pixmap.fill(self.current_color)

        # Add border
        painter = QPainter(pixmap)
        painter.setPen(Qt.black)
        painter.drawRect(0, 0, pixmap.width() - 1, pixmap.height() - 1)
        painter.end()

        self.color_btn.setIcon(QIcon(pixmap))
        self.color_btn.setStyleSheet(f"background-color: {self.current_color.name()};")

    def _load_object_classes(self):
        """Load object classes from database and populate combo box."""
        try:
            classes = self.db_manager.get_object_classes()

            # Clear and repopulate combo box
            self.class_combo.clear()
            self.class_combo.addItem("-- Select Class / Show All --", None)

            for cls in classes:
                self.class_combo.addItem(cls["class_name"], cls)

            logger.debug(f"Loaded {len(classes)} object classes")

        except Exception as e:
            logger.error(f"Error loading object classes: {e}")
            QMessageBox.warning(
                self, "Error", f"Failed to load object classes:\n{str(e)}"
            )

    def _on_polyline_toggle(self, checked: bool):
        """Handle polyline tool enable/disable."""
        self.polyline_enabled = checked

        if checked:
            self.polyline_toggle_btn.setText("Stop Drawing Polyline")
            self.polyline_toggle_btn.setStyleSheet(
                "QPushButton { background-color: #4CAF50; }"
            )
        else:
            self.polyline_toggle_btn.setText("Start Drawing Polyline")
            self.polyline_toggle_btn.setStyleSheet("")

        self.polyline_enabled_changed.emit(self.polyline_enabled)
        logger.debug(f"Polyline tool {'enabled' if checked else 'disabled'}")

    def _on_polyline_pen_width_changed(self, width: int):
        """Handle polyline pen width changes."""
        self.polyline_pen_width_changed.emit(width)
        logger.debug(f"Polyline pen width changed to {width}")

    def _on_simplify_toggle(self, state: int):
        """Handle simplify-on-finish checkbox toggle."""
        enabled = bool(state)
        self.simplify_on_finish_changed.emit(enabled)
        logger.debug(f"Simplify on finish set to {enabled}")

    def _on_eps_change(self, val: float):
        """Handle epsilon (RDP tolerance) value changes."""
        epsilon = float(val)
        self.simplify_epsilon_changed.emit(epsilon)
        logger.debug(f"Simplification epsilon changed to {epsilon}")

    def _on_show_bboxes_toggle(self, state: int):
        """Handle 'Show bounding boxes' checkbox toggle."""
        show = bool(state)
        self.show_bboxes_changed.emit(show)
        logger.debug(f"Show bounding boxes set to {show}")

    def _on_color_picker(self):
        """Open color picker dialog and update the selected object's class color."""
        if not self.current_class:
            QMessageBox.warning(
                self,
                "No Class Selected",
                "Please select an object class before changing its color.",
            )
            return

        # Use current class color (without alpha) as the base
        base_color = QColor(self.current_class.get("color", self.current_color.name()))
        color = QColorDialog.getColor(
            base_color,
            self,
            "Select Class Color",
            QColorDialog.ShowAlphaChannel,  # Allow alpha in UI, but store RGB in DB
        )

        if not color.isValid():
            return

        # Normalize to opaque RGB for storage
        new_color = QColor(color)
        new_color.setAlpha(255)
        hex_color = new_color.name()

        try:
            # Update in database
            self.db_manager.update_object_class(
                class_id=self.current_class["id"], color=hex_color
            )
        except Exception as e:
            logger.error(f"Failed to update class color in database: {e}")
            QMessageBox.critical(
                self,
                "Error",
                f"Failed to update class color in database:\n{str(e)}",
            )
            return

        # Update local class data and combo box item data
        self.current_class["color"] = hex_color
        current_index = self.class_combo.currentIndex()
        if current_index >= 0:
            self.class_combo.setItemData(current_index, dict(self.current_class))

        # Update info label text
        info_text = f"Class: {self.current_class['class_name']}\nColor: {hex_color}"
        if self.current_class.get("description"):
            info_text += f"\nDescription: {self.current_class['description']}"
        self.class_info_label.setText(info_text)

        # Use semi-transparent version for polyline pen / button preview
        class_color = QColor(hex_color)
        class_color.setAlpha(128)
        self.current_color = class_color
        self._update_color_button()
        self.polyline_pen_color_changed.emit(class_color)

        logger.debug(
            f"Updated class '{self.current_class['class_name']}' color to "
            f"{hex_color} (polyline pen alpha={class_color.alpha()})"
        )

        # Notify listeners (e.g., AnnotationTab) so they can reload/redraw
        self.class_color_changed.emit()

    def _on_class_selected(self, index: int):
        """Handle object class selection (including '-- Select Class --')."""
        class_data = self.class_combo.currentData()

        if class_data:
            self.current_class = class_data

            # Update info label
            info_text = (
                f"Class: {class_data['class_name']}\n" f"Color: {class_data['color']}"
            )
            if class_data.get("description"):
                info_text += f"\nDescription: {class_data['description']}"

            self.class_info_label.setText(info_text)

            # Update polyline pen color to match class color with semi-transparency
            class_color = QColor(class_data["color"])
            if class_color.isValid():
                # Add 50% alpha for semi-transparency
                class_color.setAlpha(128)
                self.current_color = class_color
                self._update_color_button()
                self.polyline_pen_color_changed.emit(class_color)

            self.class_selected.emit(class_data)
            logger.debug(f"Selected class: {class_data['class_name']}")
        else:
            # "-- Select Class --" chosen: clear current class and show all annotations
            self.current_class = None
            self.class_info_label.setText("No class selected")
            self.class_selected.emit(None)
            logger.debug("Class selection cleared: showing annotations for all classes")

    def _on_add_class(self):
        """Handle adding a new object class."""
        # Get class name
        class_name, ok = QInputDialog.getText(
            self, "Add Object Class", "Enter class name:"
        )

        if not ok or not class_name.strip():
            return

        class_name = class_name.strip()

        # Check if class already exists
        existing = self.db_manager.get_object_class_by_name(class_name)
        if existing:
            QMessageBox.warning(
                self, "Class Exists", f"A class named '{class_name}' already exists."
            )
            return

        # Get color
        color = QColorDialog.getColor(self.current_color, self, "Select Class Color")

        if not color.isValid():
            return

        # Get optional description
        description, ok = QInputDialog.getText(
            self, "Class Description", "Enter class description (optional):"
        )

        if not ok:
            description = None

        # Add to database
        try:
            class_id = self.db_manager.add_object_class(
                class_name, color.name(), description.strip() if description else None
            )

            logger.info(f"Added new object class: {class_name} (ID: {class_id})")

            # Reload classes and select the new one
            self._load_object_classes()

            # Find and select the newly added class
            for i in range(self.class_combo.count()):
                class_data = self.class_combo.itemData(i)
                if class_data and class_data.get("id") == class_id:
                    self.class_combo.setCurrentIndex(i)
                    break

            QMessageBox.information(
                self, "Success", f"Class '{class_name}' added successfully!"
            )

        except Exception as e:
            logger.error(f"Error adding object class: {e}")
            QMessageBox.critical(
                self, "Error", f"Failed to add object class:\n{str(e)}"
            )

    def _on_clear_annotations(self):
        """Handle clear annotations button."""
        reply = QMessageBox.question(
            self,
            "Clear Annotations",
            "Are you sure you want to clear all annotations?",
            QMessageBox.Yes | QMessageBox.No,
            QMessageBox.No,
        )

        if reply == QMessageBox.Yes:
            self.clear_annotations_requested.emit()
            logger.debug("Clear annotations requested")

    def get_current_class(self) -> Optional[Dict]:
        """Get currently selected object class."""
        return self.current_class

    def get_polyline_pen_color(self) -> QColor:
        """Get current polyline pen color."""
        return self.current_color

    def get_polyline_pen_width(self) -> int:
        """Get current polyline pen width."""
        return self.polyline_pen_width_spin.value()

    def is_polyline_enabled(self) -> bool:
        """Check if polyline tool is enabled."""
        return self.polyline_enabled