import os
import datetime
import re
from PySide6.QtWidgets import QSpacerItem, QSizePolicy, QPushButton
from PySide6.QtGui import QPixmap, Qt
from PySide6.QtWidgets import QWidget
from typing import Dict
from airunner.gui.widgets.base_widget import BaseWidget
from airunner.gui.widgets.canvas.templates.batch_container_ui import (
    Ui_batch_conatiner,
)
from airunner.gui.widgets.canvas.templates.image_layer_item_ui import (
    Ui_image_layer_item,
)
from airunner.enums import SignalCode
from airunner.utils.image.export_image import get_today_folder


class BatchContainer(BaseWidget):
    widget_class_ = Ui_batch_conatiner

    def __init__(self, *args, **kwargs):
        self.initialized = False
        self.signal_handlers = {
            SignalCode.SD_UPDATE_BATCH_IMAGES_SIGNAL: self.update_batch_images,
        }
        self.current_date_folder = (
            None  # Tracks the currently selected date folder
        )
        self.current_batch_folder = (
            None  # Tracks the currently selected batch folder
        )
        self.back_button = (
            None  # Reference to the back button when in a batch folder
        )
        super().__init__(*args, **kwargs)
        self.setup_ui_connections()

    def setup_ui_connections(self):
        """Set up UI signal connections."""
        self.ui.image_folders.currentTextChanged.connect(
            self.on_image_folders_currentTextChanged
        )

    def _add_image_layer_item(self, image_path: str, total: int, layout):
        pixmap = QPixmap(image_path)
        pixmap = pixmap.scaled(
            256,
            256,
            aspectMode=Qt.KeepAspectRatio,
            mode=Qt.SmoothTransformation,
        )
        image_layer_item = ImageLayerItemWidget(image_path)
        image_layer_item.ui.image.setPixmap(pixmap)

        if total is not None:
            image_layer_item.ui.total_images.setText(f"{total} in batch")
            image_layer_item.ui.total_images.setVisible(True)
            # Set up batch folder information for click handling
            image_layer_item.batch_folder = os.path.dirname(image_path)
            image_layer_item.is_batch = True
            image_layer_item.parent_widget = self
        else:
            # Hide the label for individual images (not part of a batch)
            image_layer_item.ui.total_images.setVisible(False)
            image_layer_item.is_batch = False

        layout.addWidget(image_layer_item)

    def _clear_layout(self, layout):
        for i in reversed(range(layout.count())):
            item = layout.itemAt(i)
            widget = item.widget() if item is not None else None
            if widget is not None:
                widget.setParent(None)
            else:
                layout.takeAt(i)

    def showEvent(self, event):
        super().showEvent(event)
        if not self.initialized:
            self.initialized = True
            self.populate_date_folders()
            self.populate_current_folder()

    def populate_date_folders(self):
        """Populate the combo box with available date folders."""
        base_path = self.path_settings.image_path
        today = datetime.datetime.now().strftime("%Y%m%d")

        # Disconnect signal temporarily to prevent triggering while populating
        self.ui.image_folders.blockSignals(True)
        self.ui.image_folders.clear()

        # Get all folders in the base path
        if os.path.exists(base_path):
            folders = [
                d
                for d in os.listdir(base_path)
                if os.path.isdir(os.path.join(base_path, d))
                and len(d) == 8
                and d.isdigit()
            ]

            # Convert to display format and sort in reverse (newest first)
            display_folders = []
            for folder in folders:
                try:
                    year = folder[0:4]
                    month = folder[4:6]
                    day = folder[6:8]
                    display_format = f"{year}-{month}-{day}"
                    display_folders.append((display_format, folder))
                except (ValueError, IndexError):
                    # Skip folders that don't match the expected format
                    pass

            # Sort by actual date (newest first)
            display_folders.sort(key=lambda x: x[1], reverse=True)

            # Add to combo box
            for display_folder, _ in display_folders:
                self.ui.image_folders.addItem(display_folder)

            # Select today's date if it exists
            today_display = f"{today[0:4]}-{today[4:6]}-{today[6:8]}"
            index = self.ui.image_folders.findText(today_display)
            if index >= 0:
                self.ui.image_folders.setCurrentIndex(index)
            elif self.ui.image_folders.count() > 0:
                # Select first item if today doesn't exist
                self.ui.image_folders.setCurrentIndex(0)

        # Reconnect signal
        self.ui.image_folders.blockSignals(False)

        # Store the currently selected date folder
        if self.ui.image_folders.currentText():
            self.current_date_folder = self._get_date_folder_path(
                self.ui.image_folders.currentText()
            )
        else:
            # If no folders exist, use today's folder
            self.current_date_folder = get_today_folder(
                self.path_settings.image_path
            )

    def _get_date_folder_path(self, display_date: str) -> str:
        """Convert display date (YYYY-MM-DD) to folder path."""
        try:
            # Convert from display format to folder name format
            date_parts = display_date.split("-")
            folder_name = f"{date_parts[0]}{date_parts[1]}{date_parts[2]}"
            return os.path.join(self.path_settings.image_path, folder_name)
        except (ValueError, IndexError):
            # Default to today's folder if conversion fails
            return get_today_folder(self.path_settings.image_path)

    def on_image_folders_currentTextChanged(self, text: str):
        """Handle selection change in the image_folders combobox."""
        if not text:
            return

        # Update current date folder and reset batch folder
        self.current_date_folder = self._get_date_folder_path(text)
        self.current_batch_folder = None

        # Remove back button if it exists
        self._remove_back_button()

        # Populate the layout with the selected date folder content
        self.populate_current_folder()

    def _create_back_button(self):
        """Create and add a back button below the combobox."""
        if self.back_button is None:
            self.back_button = QPushButton("← Back to date folder")
            self.back_button.clicked.connect(self.on_back_button_clicked)
            self.ui.gridLayout.addWidget(self.back_button, 1, 0, 1, 1)
            # Move scroll area down one row
            self.ui.gridLayout.removeWidget(self.ui.scrollArea)
            self.ui.gridLayout.addWidget(self.ui.scrollArea, 2, 0, 1, 1)

    def _remove_back_button(self):
        """Remove the back button if it exists."""
        if self.back_button is not None:
            self.ui.gridLayout.removeWidget(self.back_button)
            self.back_button.deleteLater()
            self.back_button = None

            # Move scroll area back up
            self.ui.gridLayout.removeWidget(self.ui.scrollArea)
            self.ui.gridLayout.addWidget(self.ui.scrollArea, 1, 0, 1, 1)

    def on_back_button_clicked(self):
        """Handle back button click to return to date folder view."""
        self.current_batch_folder = None
        self._remove_back_button()
        self.populate_current_folder()

    def on_batch_clicked(self, batch_folder: str):
        """Handle batch folder click to navigate into the batch."""
        self.current_batch_folder = batch_folder
        self._create_back_button()
        self.populate_current_folder()

    def find_date_loose_images(self, folder_path: str) -> list:
        """Return a list of loose images (not in batches) for the given folder."""
        if not os.path.exists(folder_path):
            return []

        loose_images = sorted(
            [
                os.path.join(folder_path, f)
                for f in os.listdir(folder_path)
                if os.path.isfile(os.path.join(folder_path, f))
            ]
        )
        return loose_images

    def find_date_batches(self, folder_path: str) -> list:
        """Return a list of batch folders and their images for the given folder."""
        if not os.path.exists(folder_path):
            return []

        batch_folders = sorted(
            [
                os.path.join(folder_path, d)
                for d in os.listdir(folder_path)
                if os.path.isdir(os.path.join(folder_path, d))
                and d.startswith("batch_")
            ],
            key=lambda x: natural_sort_key(os.path.basename(x)),
        )
        batches = []
        for batch in batch_folders:
            images = sorted(
                [
                    os.path.join(batch, f)
                    for f in os.listdir(batch)
                    if os.path.isfile(os.path.join(batch, f))
                ]
            )
            batches.append({"batch_folder": batch, "images": images})
        return batches

    def find_batch_images(self, batch_folder: str) -> list:
        """Return a list of images in a specific batch folder."""
        if not os.path.exists(batch_folder):
            return []

        return sorted(
            [
                os.path.join(batch_folder, f)
                for f in os.listdir(batch_folder)
                if os.path.isfile(os.path.join(batch_folder, f))
            ]
        )

    def populate_current_folder(self):
        """Populate the layout based on the current navigation state."""
        container = self.ui.scrollArea.widget()
        layout = container.layout()
        self._clear_layout(layout)

        if self.current_batch_folder:
            # We're viewing a batch folder - display all images in this batch
            images = self.find_batch_images(self.current_batch_folder)
            for image_path in images:
                self._add_image_layer_item(image_path, None, layout)
        else:
            # We're viewing a date folder - display batches and loose images
            batches = self.find_date_batches(self.current_date_folder)
            loose_images = self.find_date_loose_images(
                self.current_date_folder
            )

            # Add batch images
            for batch in batches:
                images = batch["images"]
                if not images:
                    continue
                self._add_image_layer_item(images[0], len(images), layout)

            # Add loose images
            for image_path in loose_images:
                self._add_image_layer_item(image_path, None, layout)

        # Add spacer at the bottom
        spacer = QSpacerItem(
            20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding
        )
        layout.addItem(spacer)

    def update_batch_images(self, data: Dict):
        """Update the layout with new batch images."""
        # Always refresh the UI to show new batches or images
        self.populate_current_folder()


class ImageLayerItemWidget(QWidget):
    def __init__(self, image_path=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ui = Ui_image_layer_item()
        self.ui.setupUi(self)
        self.ui.image.setFixedSize(256, 256)
        self.image_path = image_path
        self.batch_folder = None  # Will be set for batch folder items
        self.is_batch = False  # Flag to identify if this is a batch item
        self.parent_widget = None  # Reference to the parent BatchContainer

        # Set the cursor for better UX
        self.setCursor(Qt.PointingHandCursor)
        self.setAcceptDrops(
            False
        )  # This widget is a drag source, not a target
        self._drag_start_pos = None

    def mousePressEvent(self, event):
        super().mousePressEvent(event)
        if event.button() == Qt.LeftButton:
            self._drag_start_pos = event.pos()

    def mouseMoveEvent(self, event):
        if self.is_batch:
            # Don't allow drag for batch folders
            return
        if (
            self._drag_start_pos is not None
            and (event.pos() - self._drag_start_pos).manhattanLength() > 10
        ):
            from PySide6.QtGui import QDrag, QPixmap
            from PySide6.QtCore import QMimeData, QUrl

            drag = QDrag(self)
            mime_data = QMimeData()
            if self.image_path:
                mime_data.setUrls([QUrl.fromLocalFile(self.image_path)])
                mime_data.setText(self.image_path)
            drag.setMimeData(mime_data)

            # Set drag pixmap for visual feedback, centered on cursor
            pixmap = self.ui.image.pixmap()
            if pixmap is not None:
                scaled_pixmap = pixmap.scaled(
                    128, 128, Qt.KeepAspectRatio, Qt.SmoothTransformation
                )
                drag.setPixmap(scaled_pixmap)
                center = scaled_pixmap.rect().center()
                drag.setHotSpot(center)

            drag.exec(Qt.CopyAction)
        super().mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        super().mouseReleaseEvent(event)
        # If this is a batch item and we have a parent widget reference, navigate into the batch
        if self.is_batch and self.parent_widget and self.batch_folder:
            self.parent_widget.on_batch_clicked(self.batch_folder)
        self._drag_start_pos = None


def natural_sort_key(s):
    """Helper for natural sorting (e.g., batch_2 before batch_10)."""
    return [
        int(text) if text.isdigit() else text.lower()
        for text in re.split(r"(\d+)", s)
    ]
