import os
import shutil
from PyQt5.QtWidgets import (
    QMainWindow, QVBoxLayout, QHBoxLayout, QPushButton, QLineEdit, QListWidget,
    QFileDialog, QGraphicsView, QGraphicsScene, QLabel, QWidget, QMessageBox,
    QGroupBox, QInputDialog, QRadioButton, QCheckBox, QCompleter, QDesktopWidget
)
from PyQt5.QtGui import QPixmap, QImage, QPainter
from PyQt5.QtCore import Qt, QRectF, QStringListModel
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
IMAGES_DIR = os.path.join(BASE_DIR, "images")


class ClassificationTool(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowFlags(Qt.Window | Qt.WindowMinimizeButtonHint | Qt.CustomizeWindowHint)  # Disable close button
        self.setWindowTitle('Classification Tool')
        screen = QDesktopWidget().screenGeometry()
        self.setFixedSize(screen.width(), screen.height())
        self.showMaximized()
        self.setFixedSize(self.size())
        self.initialize_data()
        self.setup_ui()
        self.setup_stylesheet()
        self.setup_signals()

    def initialize_data(self):
        """Initialize data structures and default settings."""
        self.image_files = []
        self.current_image_index = -1
        self.image_labels = {}  # {image_path: [labels]}
        self.class_to_id = {}
        self.id_to_class = {}
        self.dataset_directory = None
        self.label_mode = "Single"
        self.default_label = None
        self.zoom_factor = 1.0
        self.min_zoom = 0.1
        self.max_zoom = 10.0
        self.undo_stack = []  # List of (image_path, labels) tuples
        self.previous_labels = []  # Store all unique labels for auto-suggestion
        self.max_image_size = 100 * 1024 * 1024  # 100 MB limit

    def setup_ui(self):
        main_layout = QHBoxLayout()
        left_panel = QVBoxLayout()
        self.back_button = QPushButton('Back')  # New Back button
        self.back_button.setToolTip("Return to Data Annotation Options")
        left_panel.addWidget(self.back_button)
        self.folder_selector_button = QPushButton('Select Image Folder')
        self.folder_selector_button.setToolTip("Select a folder containing images (.png, .jpg, .jpeg, .bmp, .tiff)")
        left_panel.addWidget(self.folder_selector_button)

        self.dataset_directory_button = QPushButton('Create Dataset Folder')
        self.dataset_directory_button.setToolTip("Create a dataset folder to store labeled images")
        left_panel.addWidget(self.dataset_directory_button)

        label_mode_group = QGroupBox("Label Mode")
        label_mode_layout = QHBoxLayout()
        self.single_label_radio = QRadioButton('Single')
        self.single_label_radio.setChecked(True)
        self.single_label_radio.setToolTip("Assign one label per image")
        label_mode_layout.addWidget(self.single_label_radio)
        self.multi_label_radio = QRadioButton('Multi')
        self.multi_label_radio.setToolTip("Assign multiple labels per image")
        label_mode_layout.addWidget(self.multi_label_radio)
        label_mode_group.setLayout(label_mode_layout)
        left_panel.addWidget(label_mode_group)

        label_group = QGroupBox("Label Assignment")
        label_layout = QVBoxLayout()
        self.label_input = QLineEdit()
        self.label_input.setPlaceholderText("Enter label")
        self.label_input.setToolTip("Enter a label and press Enter to assign it")
        self.completer = QCompleter()
        self.completer.setCaseSensitivity(Qt.CaseInsensitive)
        self.completer.setCompletionMode(QCompleter.PopupCompletion)  # Ensure popup shows reliably
        self.label_input.setCompleter(self.completer)
        label_layout.addWidget(self.label_input)
        self.add_label_button = QPushButton("Add Label")
        self.add_label_button.setToolTip("Click to assign the entered label")
        label_layout.addWidget(self.add_label_button)
        self.clear_labels_button = QPushButton("Clear All Labels")
        self.clear_labels_button.setToolTip("Remove all labels from the current image")
        label_layout.addWidget(self.clear_labels_button)
        label_group.setLayout(label_layout)
        left_panel.addWidget(label_group)

        self.previous_image_button = QPushButton('Previous Image')
        self.previous_image_button.setToolTip("Go to previous image (A key)")
        left_panel.addWidget(self.previous_image_button)

        self.next_image_button = QPushButton('Next Image')
        self.next_image_button.setToolTip("Go to next image (D key)")
        left_panel.addWidget(self.next_image_button)

        self.zoom_in_button = QPushButton('Zoom In')
        self.zoom_in_button.setToolTip("Zoom in on the image")
        left_panel.addWidget(self.zoom_in_button)

        self.zoom_out_button = QPushButton('Zoom Out')
        self.zoom_out_button.setToolTip("Zoom out on the image")
        left_panel.addWidget(self.zoom_out_button)

        self.reset_zoom_button = QPushButton('Reset Zoom')
        self.reset_zoom_button.setToolTip("Reset image zoom to default")
        left_panel.addWidget(self.reset_zoom_button)

        self.undo_button = QPushButton('Undo')
        self.undo_button.setToolTip("Undo the last label change (Ctrl+Z)")
        left_panel.addWidget(self.undo_button)

        left_panel_widget = QWidget()
        left_panel_widget.setLayout(left_panel)
        left_panel_widget.setFixedWidth(300)
        main_layout.addWidget(left_panel_widget)

        self.image_viewer = QGraphicsView()
        self.image_scene = QGraphicsScene()
        self.image_viewer.setScene(self.image_scene)
        self.image_viewer.setRenderHint(QPainter.Antialiasing)
        self.image_viewer.setRenderHint(QPainter.SmoothPixmapTransform)
        self.image_viewer.setViewportUpdateMode(QGraphicsView.FullViewportUpdate)
        self.image_viewer.setOptimizationFlag(QGraphicsView.DontAdjustForAntialiasing, True)
        main_layout.addWidget(self.image_viewer)

        right_panel = QVBoxLayout()

        default_label_group = QGroupBox("Default Label")
        default_label_layout = QVBoxLayout()
        self.default_label_checkbox = QCheckBox('Use Default Label')
        self.default_label_checkbox.setToolTip("Toggle to apply a default label (Space key)")
        default_label_layout.addWidget(self.default_label_checkbox)
        self.default_label_textbox = QLineEdit()
        self.default_label_textbox.setPlaceholderText("Enter default label")
        self.default_label_textbox.setToolTip("Enter the default label to apply automatically")
        default_label_layout.addWidget(self.default_label_textbox)
        default_label_group.setLayout(default_label_layout)
        right_panel.addWidget(default_label_group)

        label_reflector_group = QGroupBox("Assigned Labels")
        label_reflector_layout = QVBoxLayout()
        self.label_reflector = QListWidget()
        self.label_reflector.setToolTip("Select labels to delete in Multi mode (Delete key)")
        label_reflector_layout.addWidget(self.label_reflector)
        label_reflector_group.setLayout(label_reflector_layout)
        right_panel.addWidget(label_reflector_group)

        file_list_group = QGroupBox("File List")
        file_list_layout = QHBoxLayout()
        self.image_position_label = QLabel("0/0")
        self.image_position_label.setStyleSheet("color: #000000; font-size: 16px; margin-left: 10px;")
        file_list_layout.addWidget(self.image_position_label, alignment=Qt.AlignRight)
        self.file_list = QListWidget()
        self.file_list.setToolTip("Select an image to view and label")
        file_list_layout.addWidget(self.file_list)
        file_list_group.setLayout(file_list_layout)
        right_panel.addWidget(file_list_group)

        right_panel_widget = QWidget()
        right_panel_widget.setLayout(right_panel)
        right_panel_widget.setFixedWidth(300)
        main_layout.addWidget(right_panel_widget)

        central_widget = QWidget()
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)
        self.setFocusPolicy(Qt.StrongFocus)
        self.setFocus()

    def setup_stylesheet(self):
        """Apply the stylesheet to the application."""
        self.setStyleSheet("""
            QMainWindow {
                background-color: #FFFFFF;
                color: #000000;
            }
            QPushButton {
                background-color: #000000;
                color: #FFFFFF;
                border: 2px solid #FFFFFF;
                padding: 10px;
                margin: 5px;
                border-radius: 5px;
                font: 25px 'Arial';
            }
            QPushButton:hover {
                background-color: #CCCCCC;
                color: #000000;
            }
            QPushButton:pressed {
                background-color: #AAAAAA;
                color: #000000;
            }
            QLineEdit {
                background-color: #000000;
                color: #FFFFFF;
                border: 2px solid #FFFFFF;
                padding: 5px;
                margin: 5px;
                border-radius: 5px;
                font: 20px 'Arial';
            }
            QLabel {
                color: #000000;
                margin: 5px;
                font: 20px 'Arial';
            }
            QListWidget {
                background-color: #FFFFFF;
                color: #000000;
                border: 2px solid #000000;
                margin: 5px;
                border-radius: 5px;
                font: 20px 'Arial';
            }
            QGraphicsView {
                border: 2px solid #000000;
            }
            QGroupBox {
                border: 2px solid #000000;
                border-radius: 5px;
                margin-top: 10px;
                font: bold 20px 'Arial';
                color: #000000;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                subcontrol-position: top center;
                padding: 0 3px;
            }
            QRadioButton {
                color: #000000;
                font: 20px 'Arial';
                margin: 5px;
            }
            QRadioButton::indicator {
                width: 20px;
                height: 20px;
            }
            QCheckBox {
                color: #000000;
                font: 20px 'Arial';
                margin: 5px;
            }
        """)

    def setup_signals(self):
        self.back_button.clicked.connect(self.open_data_annotation_window)  # Connect Back button
        self.folder_selector_button.clicked.connect(self.select_image_folder)
        self.dataset_directory_button.clicked.connect(self.create_dataset_directory)
        self.single_label_radio.toggled.connect(self.set_label_mode)
        self.multi_label_radio.toggled.connect(self.set_label_mode)
        self.label_input.returnPressed.connect(self.add_label)
        self.label_input.textChanged.connect(self.update_completer)
        self.completer.activated[str].connect(self.on_completer_activated)
        self.add_label_button.clicked.connect(self.add_label)
        self.clear_labels_button.clicked.connect(self.clear_all_labels)
        self.previous_image_button.clicked.connect(self.previous_image_button_clicked)
        self.next_image_button.clicked.connect(self.next_image)
        self.zoom_in_button.clicked.connect(self.zoom_in)
        self.zoom_out_button.clicked.connect(self.zoom_out)
        self.reset_zoom_button.clicked.connect(self.reset_zoom)
        self.undo_button.clicked.connect(self.undo)
        self.default_label_checkbox.stateChanged.connect(self.toggle_default_label)
        self.default_label_textbox.textChanged.connect(self.update_default_label)
        self.file_list.itemSelectionChanged.connect(self.on_file_selected)
        self.label_reflector.itemSelectionChanged.connect(self.delete_selected_label)

    def open_data_annotation_window(self):
        try:
            from .fun_pyqt5 import DataAnnotationWindow
            from PyQt5.QtGui import QPixmap
            import os
            bg_image_path = os.path.join(IMAGES_DIR, 'bg_help_image.png')
            if not os.path.exists(bg_image_path):
                raise FileNotFoundError(f"Background image {bg_image_path} not found")
            self.data_annotation_window = DataAnnotationWindow(QPixmap(bg_image_path))
            self.data_annotation_window.showMaximized()
            self.hide()  # Changed from close() to hide() to preserve state
        except FileNotFoundError as e:
            QMessageBox.critical(self, "Error", str(e))
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to open Data Annotation Options: {str(e)}")

    def on_completer_activated(self, text):
        """Handle selection of a completer suggestion."""
        self.label_input.setText(text)

    def update_completer(self, text):
        """Update the auto-suggestion completer based on input text."""
        try:
            if text.strip():
                suggestions = [label for label in self.previous_labels if label.lower().startswith(text.lower())]
                self.completer.setModel(QStringListModel(suggestions))
                self.completer.complete()
            else:
                self.completer.setModel(QStringListModel(self.previous_labels))
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to update completer: {str(e)}")

    def set_label_mode(self):
        """Set the label mode and update selection behavior."""
        try:
            if self.single_label_radio.isChecked():
                self.label_mode = "Single"
                self.label_reflector.setSelectionMode(QListWidget.SingleSelection)
            elif self.multi_label_radio.isChecked():
                self.label_mode = "Multi"
                self.label_reflector.setSelectionMode(QListWidget.MultiSelection)
            self.update_label_reflector()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to set label mode: {str(e)}")

    def toggle_default_label(self, state):
        """Toggle the default label and apply it if enabled."""
        try:
            if state == Qt.Checked and self.default_label_textbox.text().strip():
                self.default_label = self.default_label_textbox.text().strip()
                self.apply_default_label()
            else:
                self.default_label = None
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to toggle default label: {str(e)}")

    def update_default_label(self, text):
        """Update the default label and apply it if the checkbox is checked."""
        try:
            if self.default_label_checkbox.isChecked() and text.strip():
                self.default_label = text.strip()
                self.apply_default_label()
            else:
                self.default_label = None
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to update default label: {str(e)}")

    def apply_default_label(self):
        """Apply the default label to the current image, respecting mode."""
        try:
            if self.default_label and self.current_image_index >= 0:
                image_path = self.image_files[self.current_image_index]
                self.undo_stack.append((image_path, self.image_labels.get(image_path, [])[:]))
                if image_path in self.image_labels and self.image_labels[image_path]:
                    if self.label_mode == "Single":
                        reply = QMessageBox.question(self, "Overwrite Labels",
                                                    "This image already has labels. Replace with default label?",
                                                    QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
                        if reply == QMessageBox.No:
                            return
                        for label in self.image_labels[image_path]:
                            label_dir = os.path.join(self.dataset_directory, label)
                            dest_path = os.path.join(label_dir, os.path.basename(image_path))
                            if os.path.exists(dest_path):
                                self.delete_file(dest_path)
                        self.image_labels[image_path] = [self.default_label]
                    else:
                        if self.default_label not in self.image_labels[image_path]:
                            self.image_labels[image_path].append(self.default_label)
                else:
                    self.image_labels[image_path] = [self.default_label]
                self.get_class_id(self.default_label)
                self.save_label(image_path, self.default_label)
                self.update_label_reflector()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to apply default label: {str(e)}")

    def prompt_label(self):
        """Prompt for a label, using default if applicable."""
        try:
            if self.default_label_checkbox.isChecked() and self.default_label_textbox.text().strip():
                return self.default_label_textbox.text().strip()
            else:
                label = self.label_input.text().strip()
                if not label:
                    return None
                return label
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error prompting label: {str(e)}")
            return None

    def select_image_folder(self):
        """Select a folder and load valid images (limited to 100)."""
        folder = QFileDialog.getExistingDirectory(self, "Select Image Folder")
        if folder:
            try:
                if not os.access(folder, os.R_OK):
                    raise PermissionError("No read permission for the selected folder")
                self.image_files = []
                valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.tiff')
                all_files = [os.path.join(folder, f) for f in os.listdir(folder) if f.lower().endswith(valid_extensions)]
                for file_path in all_files[:100]:
                    file_size = os.path.getsize(file_path)
                    if file_size > self.max_image_size:
                        continue
                    try:
                        with open(file_path, 'rb'):
                            image = QImage(file_path)
                            if image.isNull() or image.width() == 0 or image.height() == 0:
                                continue
                        self.image_files.append(file_path)
                    except Exception:
                        continue
                self.file_list.blockSignals(True)
                self.file_list.clear()
                self.file_list.addItems([os.path.basename(f) for f in self.image_files])
                self.file_list.blockSignals(False)
                if self.image_files:
                    self.current_image_index = 0
                    self.file_list.setCurrentRow(0)
                    self.display_image(self.image_files[0])
                    self.update_image_position_label()
                    self.update_navigation_buttons()
                    if len(all_files) > 100:
                        QMessageBox.information(self, "Info", f"Loaded first 100 valid images. Total: {len(all_files)}")
                else:
                    raise ValueError("No valid image files found")
            except PermissionError as e:
                QMessageBox.warning(self, "Permission Error", str(e))
            except ValueError as e:
                QMessageBox.warning(self, "Invalid Folder", str(e))
            except Exception as e:
                QMessageBox.warning(self, "Error", f"Unexpected error: {str(e)}")

    def create_dataset_directory(self):
        """Create a dataset directory for storing labeled images."""
        try:
            directory = QFileDialog.getExistingDirectory(self, "Select Parent Directory for Dataset")
            if directory:
                dataset_name, ok = QInputDialog.getText(self, "Dataset Name", "Enter dataset folder name:")
                if ok and dataset_name:
                    self.dataset_directory = os.path.join(directory, dataset_name)
                    os.makedirs(self.dataset_directory, exist_ok=True)
                    self.load_classes()
                    self.load_labels()
                    QMessageBox.information(self, "Success", f"Dataset directory '{dataset_name}' created")
        except PermissionError as e:
            QMessageBox.warning(self, "Permission Error", f"No permission to create directory: {str(e)}")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to create dataset directory: {str(e)}")

    def load_classes(self):
        """Load class labels from classes.txt."""
        if not self.dataset_directory:
            return
        classes_file = os.path.join(self.dataset_directory, "classes.txt")
        self.class_to_id.clear()
        self.id_to_class.clear()
        try:
            if os.path.exists(classes_file):
                with open(classes_file, 'r') as f:
                    for class_id, label in enumerate(f.read().splitlines()):
                        self.class_to_id[label] = class_id
                        self.id_to_class[class_id] = label
                        if label not in self.previous_labels:
                            self.previous_labels.append(label)
                            self.update_completer(self.label_input.text())
        except PermissionError:
            QMessageBox.warning(self, "Error", f"No permission to read classes file")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to load classes: {str(e)}")

    def update_classes_file(self):
        """Update classes.txt with current class labels."""
        if not self.dataset_directory:
            return
        classes_file = os.path.join(self.dataset_directory, "classes.txt")
        content = "\n".join(self.id_to_class[class_id] for class_id in sorted(self.id_to_class.keys())) + "\n"
        self.save_file(classes_file, content)

    def get_class_id(self, label):
        """Get or assign a class ID for a label."""
        try:
            if label not in self.class_to_id:
                class_id = len(self.class_to_id)
                self.class_to_id[label] = class_id
                self.id_to_class[class_id] = label
                self.update_classes_file()
            # Always add to previous_labels for suggestion history
            if label not in self.previous_labels:
                self.previous_labels.append(label)
            self.update_completer(self.label_input.text())
            return self.class_to_id[label]
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to assign class ID: {str(e)}")
            return None

    def display_image(self, image_path):
        """Display an image in the viewer with reset zoom."""
        try:
            file_size = os.path.getsize(image_path)
            if file_size > self.max_image_size:
                raise ValueError(f"Image file is too large ({file_size} bytes)")
            image = QImage(image_path)
            if image.isNull():
                raise ValueError("Failed to load image")
            if image.width() == 0 or image.height() == 0:
                raise ValueError("Image has zero dimensions")
            pixmap = QPixmap.fromImage(image)
            del image  # Release QImage to free memory
            if pixmap.isNull():
                raise ValueError("Failed to create pixmap")
            self.image_scene.clear()
            self.image_scene.addPixmap(pixmap)
            self.image_scene.setSceneRect(QRectF(pixmap.rect()))
            del pixmap  # Release QPixmap to free memory
            self.reset_zoom()
            self.image_viewer.viewport().update()
            self.update_label_reflector()
            if self.default_label_checkbox.isChecked() and self.default_label_textbox.text().strip():
                self.apply_default_label()
        except ValueError as e:
            QMessageBox.warning(self, "Error", str(e))
            self.skip_to_next_valid_image()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to display image: {str(e)}")
            self.skip_to_next_valid_image()

    def skip_to_next_valid_image(self):
        """Skip to the next valid image if the current one fails."""
        try:
            while self.current_image_index < len(self.image_files) - 1:
                self.current_image_index += 1
                file_path = self.image_files[self.current_image_index]
                file_size = os.path.getsize(file_path)
                if file_size > self.max_image_size:
                    continue
                image = QImage(file_path)
                if image.isNull() or image.width() == 0 or image.height() == 0:
                    continue
                self.file_list.blockSignals(True)
                self.file_list.setCurrentRow(self.current_image_index)
                self.file_list.blockSignals(False)
                self.display_image(file_path)
                self.update_image_position_label()
                self.update_navigation_buttons()
                return
            self.current_image_index = -1
            self.image_scene.clear()
            self.update_label_reflector()
            self.update_image_position_label()
            self.update_navigation_buttons()
            QMessageBox.warning(self, "Error", "No valid images found in the folder")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to skip to next image: {str(e)}")

    def add_label(self):
        """Add a label to the current image."""
        try:
            if self.current_image_index < 0 or not self.dataset_directory:
                QMessageBox.warning(self, "Error", "Select an image folder and create a dataset directory first")
                return
            label = self.prompt_label()
            if not label:
                QMessageBox.warning(self, "Error", "Enter a valid label")
                return
            image_path = self.image_files[self.current_image_index]
            self.undo_stack.append((image_path, self.image_labels.get(image_path, [])[:]))
            if image_path in self.image_labels:
                if self.label_mode == "Single":
                    self.image_labels[image_path] = [label]
                else:
                    if label not in self.image_labels[image_path]:
                        self.image_labels[image_path].append(label)
            else:
                self.image_labels[image_path] = [label]
            self.get_class_id(label)
            self.save_label(image_path, label)
            self.update_label_reflector()
            self.label_input.clear()
            if self.label_mode == "Single":
                self.next_image()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to add label: {str(e)}")

    def save_label(self, image_path, label):
        """Save a label by copying the image to the label's directory."""
        try:
            label_dir = os.path.join(self.dataset_directory, label)
            os.makedirs(label_dir, exist_ok=True)
            dest_path = os.path.join(label_dir, os.path.basename(image_path))
            shutil.copy2(image_path, dest_path)
            self.save_labels_file()
        except PermissionError:
            QMessageBox.warning(self, "Error", f"No permission to save label '{label}'")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to save label: {str(e)}")

    def save_labels_file(self):
        """Save the image labels to labels.txt."""
        if not self.dataset_directory:
            return
        labels_file = os.path.join(self.dataset_directory, "labels.txt")
        try:
            content = "\n".join(
                f"{os.path.basename(image_path)} {','.join(labels)}"
                for image_path, labels in self.image_labels.items()
            ) + "\n"
            self.save_file(labels_file, content)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to save labels file: {str(e)}")

    def load_labels(self):
        """Load image labels from labels.txt."""
        if not self.dataset_directory:
            return
        labels_file = os.path.join(self.dataset_directory, "labels.txt")
        try:
            if os.path.exists(labels_file):
                with open(labels_file, 'r') as f:
                    for line in f:
                        parts = line.strip().split(' ', 1)
                        if len(parts) == 2:
                            image_name, labels = parts
                            image_path = next((ip for ip in self.image_files if os.path.basename(ip) == image_name), None)
                            if image_path:
                                self.image_labels[image_path] = labels.split(',')
                                for label in self.image_labels[image_path]:
                                    self.get_class_id(label)
        except PermissionError:
            QMessageBox.warning(self, "Error", f"No permission to read labels file")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to load labels: {str(e)}")

    def clear_all_labels(self):
        """Clear all labels from the current image."""
        try:
            if self.current_image_index < 0:
                return
            image_path = self.image_files[self.current_image_index]
            if image_path in self.image_labels:
                self.undo_stack.append((image_path, self.image_labels[image_path][:]))
                for label in self.image_labels[image_path]:
                    label_dir = os.path.join(self.dataset_directory, label)
                    dest_path = os.path.join(label_dir, os.path.basename(image_path))
                    if os.path.exists(dest_path):
                        self.delete_file(dest_path)
                del self.image_labels[image_path]
                self.save_labels_file()
                self.update_label_reflector()
                QMessageBox.information(self, "Success", "All labels cleared")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to clear labels: {str(e)}")

    def delete_selected_label(self):
        """Delete selected labels in Multi mode."""
        try:
            if self.current_image_index < 0 or self.label_mode != "Multi":
                return
            selected_items = self.label_reflector.selectedItems()
            if not selected_items:
                QMessageBox.warning(self, "Error", "No label selected")
                return
            image_path = self.image_files[self.current_image_index]
            if image_path not in self.image_labels:
                return
            self.undo_stack.append((image_path, self.image_labels[image_path][:]))
            for item in selected_items:
                label = item.text()
                if label in self.image_labels[image_path]:
                    self.image_labels[image_path].remove(label)
                    label_dir = os.path.join(self.dataset_directory, label)
                    dest_path = os.path.join(label_dir, os.path.basename(image_path))
                    if os.path.exists(dest_path):
                        self.delete_file(dest_path)
            if not self.image_labels[image_path]:
                del self.image_labels[image_path]
            self.save_labels_file()
            self.update_label_reflector()
            QMessageBox.information(self, "Success", "Selected labels deleted")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to delete labels: {str(e)}")

    def update_label_reflector(self):
        """Update the label reflector with current image labels."""
        try:
            self.label_reflector.blockSignals(True)
            self.label_reflector.clear()
            if self.current_image_index >= 0:
                image_path = self.image_files[self.current_image_index]
                if image_path in self.image_labels:
                    self.label_reflector.addItems(self.image_labels[image_path])
            self.label_reflector.blockSignals(False)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to update label reflector: {str(e)}")

    def next_image(self):
        """Navigate to the next image."""
        try:
            if self.current_image_index < len(self.image_files) - 1:
                self.current_image_index += 1
                self.file_list.blockSignals(True)
                self.file_list.setCurrentRow(self.current_image_index)
                self.file_list.blockSignals(False)
                self.display_image(self.image_files[self.current_image_index])
                self.update_image_position_label()
                self.update_navigation_buttons()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to navigate to next image: {str(e)}")

    def previous_image(self, clear_previous_labels=False):
        """Navigate to the previous image, optionally clearing labels."""
        try:
            if self.current_image_index > 0:
                self.current_image_index -= 1
                if clear_previous_labels and self.label_mode == "Single":
                    previous_image_path = self.image_files[self.current_image_index]
                    if previous_image_path in self.image_labels:
                        self.undo_stack.append((previous_image_path, self.image_labels[previous_image_path][:]))
                        for label in self.image_labels[previous_image_path]:
                            label_dir = os.path.join(self.dataset_directory, label)
                            dest_path = os.path.join(label_dir, os.path.basename(previous_image_path))
                            if os.path.exists(dest_path):
                                self.delete_file(dest_path)
                        del self.image_labels[previous_image_path]
                        self.save_labels_file()
                self.file_list.blockSignals(True)
                self.file_list.setCurrentRow(self.current_image_index)
                self.file_list.blockSignals(False)
                self.display_image(self.image_files[self.current_image_index])
                self.update_image_position_label()
                self.update_navigation_buttons()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to navigate to previous image: {str(e)}")

    def previous_image_button_clicked(self):
        """Handle previous image button click."""
        try:
            self.previous_image(clear_previous_labels=False)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to navigate to previous image: {str(e)}")

    def on_file_selected(self):
        """Handle file selection from the file list."""
        try:
            selected_items = self.file_list.selectedItems()
            if selected_items:
                selected_file = selected_items[0].text()
                self.current_image_index = next(i for i, ip in enumerate(self.image_files) if os.path.basename(ip) == selected_file)
                self.display_image(self.image_files[self.current_image_index])
                self.update_image_position_label()
                self.update_navigation_buttons()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to select file: {str(e)}")

    def update_image_position_label(self):
        """Update the image position label."""
        try:
            current = self.current_image_index + 1 if self.current_image_index >= 0 else 0
            total = len(self.image_files)
            self.image_position_label.setText(f"{current}/{total}")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to update position label: {str(e)}")

    def update_navigation_buttons(self):
        """Enable/disable navigation buttons based on position."""
        try:
            self.previous_image_button.setEnabled(self.current_image_index > 0 and len(self.image_files) > 0)
            self.next_image_button.setEnabled(self.current_image_index < len(self.image_files) - 1 and len(self.image_files) > 0)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to update navigation buttons: {str(e)}")

    def zoom_in(self):
        """Zoom in on the image with limits."""
        try:
            if self.zoom_factor < self.max_zoom:
                self.zoom_factor *= 1.1
                self.image_viewer.scale(1.1, 1.1)
                self.image_viewer.viewport().update()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to zoom in: {str(e)}")

    def zoom_out(self):
        """Zoom out on the image with limits."""
        try:
            if self.zoom_factor > self.min_zoom:
                self.zoom_factor *= 0.9
                self.image_viewer.scale(0.9, 0.9)
                self.image_viewer.viewport().update()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to zoom out: {str(e)}")

    def reset_zoom(self):
        """Reset the image zoom to default."""
        try:
            self.zoom_factor = 1.0
            self.image_viewer.resetTransform()
            self.image_viewer.fitInView(self.image_scene.sceneRect(), Qt.KeepAspectRatio)
            self.image_viewer.viewport().update()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to reset zoom: {str(e)}")

    def undo(self):
        """Undo the last label change."""
        try:
            if self.undo_stack:
                image_path, previous_labels = self.undo_stack.pop()
                if previous_labels:
                    self.image_labels[image_path] = previous_labels
                else:
                    self.image_labels.pop(image_path, None)
                for label in self.image_labels.get(image_path, []):
                    self.save_label(image_path, label)
                self.save_labels_file()
                self.update_label_reflector()
                QMessageBox.information(self, "Success", "Undo successful")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to undo: {str(e)}")

    def save_file(self, file_path, content, mode='w'):
        """Helper method to save content to a file."""
        try:
            with open(file_path, mode) as f:
                f.write(content)
        except PermissionError:
            QMessageBox.warning(self, "Error", f"Permission denied writing to {file_path}")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to save file {file_path}: {str(e)}")

    def delete_file(self, file_path):
        """Helper method to delete a file."""
        try:
            os.remove(file_path)
        except PermissionError:
            QMessageBox.warning(self, "Error", f"No permission to delete {file_path}")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to delete file {file_path}: {str(e)}")

    def keyPressEvent(self, event):
        """Handle keyboard shortcuts."""
        try:
            if event.key() == Qt.Key_D:
                self.next_image()
            elif event.key() == Qt.Key_A:
                self.previous_image(clear_previous_labels=self.label_mode == "Single")
            elif event.key() == Qt.Key_Delete and self.label_mode == "Multi":
                self.delete_selected_label()
            elif event.key() == Qt.Key_Delete and self.label_mode == "Single":
                self.clear_all_labels()
            elif event.key() == Qt.Key_Space:
                self.default_label_checkbox.toggle()
            elif event.key() == Qt.Key_Z and event.modifiers() == Qt.ControlModifier:
                self.undo()
            super().keyPressEvent(event)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Key press error: {str(e)}")



if __name__ == "__main__":
    from PyQt5.QtWidgets import QApplication
    import sys
    try:
        app = QApplication(sys.argv)
        window = ClassificationTool()
        window.show()
        sys.exit(app.exec_())
    except Exception as e:
        print(f"Application crashed: {str(e)}")
        raise