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

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

from PySide6.QtWidgets import (
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QLabel,
    QGroupBox,
    QPushButton,
    QComboBox,
    QSpinBox,
    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 pen tool
    - Color selection for pen
    - Object class selection
    - Add new object classes
    - Pen width control
    - Clear annotations

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

    pen_enabled_changed = Signal(bool)
    pen_color_changed = Signal(QColor)
    pen_width_changed = Signal(int)
    class_selected = Signal(dict)
    clear_annotations_requested = Signal()
    process_annotations_requested = Signal()
    show_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.pen_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()

        # Pen Tool Group
        pen_group = QGroupBox("Pen Tool")
        pen_layout = QVBoxLayout()

        # Enable/Disable pen
        button_layout = QHBoxLayout()
        self.pen_toggle_btn = QPushButton("Enable Pen")
        self.pen_toggle_btn.setCheckable(True)
        self.pen_toggle_btn.clicked.connect(self._on_pen_toggle)
        button_layout.addWidget(self.pen_toggle_btn)
        pen_layout.addLayout(button_layout)

        # Pen width control
        width_layout = QHBoxLayout()
        width_layout.addWidget(QLabel("Pen Width:"))
        self.pen_width_spin = QSpinBox()
        self.pen_width_spin.setMinimum(1)
        self.pen_width_spin.setMaximum(20)
        self.pen_width_spin.setValue(3)
        self.pen_width_spin.valueChanged.connect(self._on_pen_width_changed)
        width_layout.addWidget(self.pen_width_spin)
        width_layout.addStretch()
        pen_layout.addLayout(width_layout)

        # Color selection
        color_layout = QHBoxLayout()
        color_layout.addWidget(QLabel("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()
        pen_layout.addLayout(color_layout)

        pen_group.setLayout(pen_layout)
        layout.addWidget(pen_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 class button
        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)

        # 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()

        self.process_btn = QPushButton("Process Annotations")
        self.process_btn.clicked.connect(self._on_process_annotations)
        self.process_btn.setStyleSheet(
            "QPushButton { background-color: #2196F3; color: white; font-weight: bold; }"
        )
        actions_layout.addWidget(self.process_btn)

        self.show_btn = QPushButton("Show Saved Annotations")
        self.show_btn.clicked.connect(self._on_show_annotations)
        self.show_btn.setStyleSheet(
            "QPushButton { background-color: #4CAF50; color: white; }"
        )
        actions_layout.addWidget(self.show_btn)

        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 --", 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_pen_toggle(self, checked: bool):
        """Handle pen tool enable/disable."""
        self.pen_enabled = checked

        if checked:
            self.pen_toggle_btn.setText("Disable Pen")
            self.pen_toggle_btn.setStyleSheet(
                "QPushButton { background-color: #4CAF50; }"
            )
        else:
            self.pen_toggle_btn.setText("Enable Pen")
            self.pen_toggle_btn.setStyleSheet("")

        self.pen_enabled_changed.emit(self.pen_enabled)
        logger.debug(f"Pen tool {'enabled' if checked else 'disabled'}")

    def _on_pen_width_changed(self, width: int):
        """Handle pen width changes."""
        self.pen_width_changed.emit(width)
        logger.debug(f"Pen width changed to {width}")

    def _on_color_picker(self):
        """Open color picker dialog with alpha support."""
        color = QColorDialog.getColor(
            self.current_color,
            self,
            "Select Pen Color",
            QColorDialog.ShowAlphaChannel,  # Enable alpha channel selection
        )

        if color.isValid():
            self.current_color = color
            self._update_color_button()
            self.pen_color_changed.emit(color)
            logger.debug(
                f"Pen color changed to {color.name()} with alpha {color.alpha()}"
            )

    def _on_class_selected(self, index: int):
        """Handle object class selection."""
        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 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.pen_color_changed.emit(class_color)

            self.class_selected.emit(class_data)
            logger.debug(f"Selected class: {class_data['class_name']}")
        else:
            self.current_class = None
            self.class_info_label.setText("No class selected")

    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 _on_process_annotations(self):
        """Handle process annotations button."""
        if not self.current_class:
            QMessageBox.warning(
                self,
                "No Class Selected",
                "Please select an object class before processing annotations.",
            )
            return

        self.process_annotations_requested.emit()
        logger.debug("Process annotations requested")

    def _on_show_annotations(self):
        """Handle show annotations button."""
        self.show_annotations_requested.emit()
        logger.debug("Show annotations requested")

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

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

    def get_pen_width(self) -> int:
        """Get current pen width."""
        return self.pen_width_spin.value()

    def is_pen_enabled(self) -> bool:
        """Check if pen tool is enabled."""
        return self.pen_enabled