import os
import sys
import random
import shutil
from PIL import Image
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QVBoxLayout, QHBoxLayout, QMessageBox
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QColor


class ImageComparison(QWidget):
    def __init__(self, image_folder, discarded_folder):
        super().__init__()
        self.image_folder = image_folder
        self.discarded_folder = discarded_folder
        self.image1_path = None
        self.image2_path = None
        self.image_labels = [QLabel(self) for _ in range(2)]
        for label in self.image_labels:
            label.setAlignment(Qt.AlignCenter)

        layout = QHBoxLayout()
        for label in self.image_labels:
            layout.addWidget(label, 1)

        self.setLayout(layout)
        self.remaining_images = None

        self.setWindowTitle('Image Comparison')
        screen_geometry = QApplication.desktop().screenGeometry()
        self.setGeometry(0, 0, screen_geometry.width(), screen_geometry.height())

        self.update_images()

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setPen(QPen(Qt.black, 2, Qt.SolidLine))
        center_x = self.width() // 2
        painter.drawLine(center_x, 0, center_x, self.height())

    def update_images(self):
        image_files = [f for f in os.listdir(self.image_folder)
                       if os.path.isfile(os.path.join(self.image_folder, f)) and f.lower().endswith('.png')]

        if len(image_files) == 2:  # Display remaining images and show popup
            self.remaining_images = image_files
            self.update_remaining_images()
            self.display_remaining_images_popup()
            return

        if len(image_files) < 2:  # Stop when there are fewer than two images remaining
            print("Finished!")
            self.close()
            return

        self.image1_path, self.image2_path = self.get_random_images(image_files)

        self.update_image(self.image1_path, self.image_labels[0])
        self.update_image(self.image2_path, self.image_labels[1])

    def update_image(self, image_path, label):
        image = Image.open(os.path.join(self.image_folder, image_path)).convert("RGB")

        available_height = self.height() - 40  # Adjust the available height by subtracting the padding (40 pixels)

        # Calculate the maximum height for the image
        max_height = available_height - 4  # Subtract 4 pixels for the dividing line and additional padding

        # Calculate the maximum width for the images
        max_width = (self.width() - 40) // 2  # Half of the window's width minus a 40-pixel safe space

        # Scale the image while maintaining aspect ratio
        image_width, image_height = image.size

        scale_factor_height = max_height / image_height
        scale_factor_width = max_width / image_width

        scale_factor = min(scale_factor_height, scale_factor_width)
        new_width = int(image_width * scale_factor)
        new_height = int(image_height * scale_factor)
        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

        qimage = self.convert_pil_to_qimage(image)
        pixmap = QPixmap.fromImage(qimage)

        # Adjust the label size to avoid overlapping with the divider line
        label.setMaximumSize(max_width, available_height)

        label.setPixmap(pixmap)


    def convert_pil_to_qimage(self, pil_image):
        data = pil_image.convert("RGBA").tobytes("raw", "RGBA")
        qimage = QImage(data, pil_image.size[0], pil_image.size[1], QImage.Format_RGBA8888)
        return qimage

    def keyPressEvent(self, event):
        key_mapping = {
            Qt.Key_Left: self.image1_path,
            Qt.Key_A: self.image1_path,
            Qt.Key_Right: self.image2_path,
            Qt.Key_D: self.image2_path,
            Qt.Key_Down: (self.image1_path, self.image2_path),
            Qt.Key_S: (self.image1_path, self.image2_path),
            Qt.Key_Up: None,
            Qt.Key_W: None,
        }

        key = event.key()
        image_paths = key_mapping.get(key)

        if image_paths:
            if isinstance(image_paths, tuple):
                for image_path in image_paths:
                    self.move_to_discarded(image_path)
            else:
                self.move_to_discarded(image_paths)
        elif key in (Qt.Key_Up, Qt.Key_W):
            pass  # Do nothing, no image will be moved

        self.update_images()

    def move_to_discarded(self, image_path):
        os.makedirs(self.discarded_folder, exist_ok=True)
        shutil.move(os.path.join(self.image_folder, image_path), self.discarded_folder)
        print(f"{image_path} discarded")

    def get_random_images(self, image_files):
        image1_path = random.choice(image_files)
        image_files.remove(image1_path)
        image2_path = random.choice(image_files)
        return image1_path, image2_path

    def update_remaining_images(self):
        for image_path, label in zip(self.remaining_images, self.image_labels):
            self.update_image(image_path, label)

    def display_remaining_images_popup(self):
        # Show the remaining images in the main window
        self.update_remaining_images()

        # Display a message box with the remaining image paths
        msg_box = QMessageBox()
        msg_box.setWindowTitle("Image Comparison")
        msg_box.setText("2 images remaining")
        msg_box.setInformativeText("Remaining images:\n\n" + "\n".join(self.remaining_images))

        # Show the message box
        msg_box.exec_()


def main():
    app = QApplication(sys.argv)

    if len(sys.argv) > 1:
        folder_path = sys.argv[1]
    else:
        folder_path = input("Enter the folder path: ")

    discarded_dir = os.path.join(folder_path, "discarded")
    image_comparison = ImageComparison(folder_path, discarded_dir)
    image_comparison.showMaximized()

    sys.exit(app.exec_())


if __name__ == '__main__':
    main()