import sys
import os
import glob
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QPushButton, QLabel,
                             QFileDialog, QComboBox, QLineEdit, QHBoxLayout, QCheckBox)
from PyQt5.QtCore import QRunnable, QThreadPool, pyqtSignal, QObject
from PIL import Image
import shutil

# 图片压缩函数
def compress_image(input_file_path, output_file_path, options):
    max_width, max_size_kb, output_format, keep_dimensions, _ = options

    with Image.open(input_file_path) as img:
        # 如果图片已符合要求，则将其复制到输出目录
        if (keep_dimensions or img.width <= max_width) and os.path.getsize(input_file_path) / 1024 <= max_size_kb:
            shutil.copyfile(input_file_path, output_file_path)
            return

        # 计算新尺寸并调整图片大小，如果不保持原尺寸
        if not keep_dimensions and img.width > max_width:
            height = int(max_width * img.height / img.width)
            img = img.resize((max_width, height), Image.Resampling.LANCZOS)

        # 确定输出格式
        if output_format == '原格式':
            img_format = img.format
            if not img_format:
                img_format = 'JPEG'  # 默认格式，如果原格式无法确定
            save_path = output_file_path
        else:
            img_format = 'JPEG' if output_format == 'JPEG' else 'PNG'
            save_ext = '.jpg' if output_format == 'JPEG' else '.png'
            save_path = os.path.splitext(file_path)[0] + save_ext


        # 转换RGBA为RGB（如果需要）
        if img.mode == 'RGBA' and img_format == 'JPEG':
            img = img.convert('RGB')

        # 保存图片，循环减少质量直到大小满足要求
        if max_size_kb:
            quality = 95  # 初始质量
            while True:
                from io import BytesIO
                temp_buffer = BytesIO()
                img.save(temp_buffer, format=img_format, quality=quality)
                size_kb = temp_buffer.tell() / 1024
                if size_kb <= max_size_kb or quality <= 10:
                    break
                quality -= 5  # 减少质量
            img.save(output_file_path, img_format, quality=quality)
        else:
            img.save(output_file_path, img_format)

class CompressTask(QRunnable):
    def __init__(self, input_path, output_path, options, callback):
        super().__init__()
        self.input_path = input_path
        self.output_path = output_path
        self.options = options
        self.callback = callback

    def run(self):
        compress_image(self.input_path, self.output_path, self.options)
        self.callback.emit(self.input_path)  # 发送信号更新UI

class WorkerSignals(QObject):
    # 定义一个用于更新UI的信号
    finished = pyqtSignal(str)


# 图片压缩类
class ImageCompressor(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.threadpool = QThreadPool()        
        self.tasksCompleted = 0  # 用于跟踪完成的任务数量
        self.totalTasks = 0  # 总任务数量
        self.source_dir_path = ''  # 初始化源目录路径为空
        self.output_dir_path = ''  # 初始化输出目录路径为空

    def initUI(self):
        self.setWindowTitle('Image Compressor')
        self.setGeometry(300, 300, 350, 250)

        layout = QVBoxLayout()

        # 格式选择
        format_layout = QHBoxLayout()
        self.formatCombo = QComboBox()
        self.formatCombo.addItems(['原格式', 'JPEG', 'PNG'])
        format_layout.addWidget(QLabel('输出格式:'))
        format_layout.addWidget(self.formatCombo)
        layout.addLayout(format_layout)

        # 尺寸输入
        size_layout = QHBoxLayout()
        self.widthInput = QLineEdit('1080')
        self.keepDimensions = QCheckBox('保持原尺寸')
        self.keepDimensions.stateChanged.connect(self.toggleWidthInput)
        size_layout.addWidget(QLabel('最大宽度:'))
        size_layout.addWidget(self.widthInput)
        size_layout.addWidget(self.keepDimensions)
        layout.addLayout(size_layout)

        # 文件大小输入
        size_kb_layout = QHBoxLayout()
        self.sizeInput = QLineEdit('500')
        self.noSizeLimit = QCheckBox('不限制大小')
        self.noSizeLimit.stateChanged.connect(self.toggleSizeInput)
        size_kb_layout.addWidget(QLabel('最大大小(KB):'))
        size_kb_layout.addWidget(self.sizeInput)
        size_kb_layout.addWidget(self.noSizeLimit)
        layout.addLayout(size_kb_layout)

        # 信息显示
        self.label = QLabel('选择要压缩图片的源目录和输出目录')
        self.label.setStyleSheet("color: red")
        layout.addWidget(self.label)

        # 源目录选择
        self.sourceDirLabel = QLabel('选择一个源目录来压缩图片')
        layout.addWidget(self.sourceDirLabel)

        self.btnSelectSourceDirectory = QPushButton('选择源目录', self)
        self.btnSelectSourceDirectory.clicked.connect(self.openSourceDirectory)
        layout.addWidget(self.btnSelectSourceDirectory)

        # 输出目录选择和显示
        self.outputDirLabel = QLabel('选择一个输出目录来保存图片')
        layout.addWidget(self.outputDirLabel)

        self.btnSelectOutputDirectory = QPushButton('选择输出目录', self)
        self.btnSelectOutputDirectory.clicked.connect(self.openOutputDirectory)
        layout.addWidget(self.btnSelectOutputDirectory)

        # 开始压缩按钮
        self.btnStartCompress = QPushButton('开始压缩', self)
        self.btnStartCompress.clicked.connect(self.startCompression)
        self.btnStartCompress.setEnabled(False)  # 初始时禁用
        layout.addWidget(self.btnStartCompress)


        self.setLayout(layout)

    def toggleWidthInput(self, state):
        self.widthInput.setEnabled(state == 0)

    def toggleSizeInput(self, state):
        self.sizeInput.setEnabled(state == 0)

    def openDirectory(self):
        self.dir_path = QFileDialog.getExistingDirectory(self, "选择目录")
        if self.dir_path:
            self.label.setText(f'选中目录: {self.dir_path}')
            self.btnStartCompress.setEnabled(True)  # 启用开始压缩按钮
            
    def openSourceDirectory(self):
        self.source_dir_path = QFileDialog.getExistingDirectory(self, "选择源目录")
        if self.source_dir_path:
            self.sourceDirLabel.setText(f'选中源目录: {self.source_dir_path}')
            self.validateDirectories()

    def openOutputDirectory(self):
        self.output_dir_path = QFileDialog.getExistingDirectory(self, "选择输出目录")
        if self.output_dir_path:
            self.outputDirLabel.setText(f'选中输出目录: {self.output_dir_path}')
            self.validateDirectories()

    def validateDirectories(self):
        if self.source_dir_path and self.output_dir_path and self.source_dir_path != self.output_dir_path:
            self.btnStartCompress.setEnabled(True)
        else:
            self.btnStartCompress.setEnabled(False)


    def startCompression(self):
        if not self.source_dir_path or not self.output_dir_path:
            self.label.setText('未选择目录，必须选择源目录及输出目录！')
            return

        # 禁用所有控件
        self.disableWidgets(True)

        # 读取压缩选项
        options = (
            int(self.widthInput.text()) if not self.keepDimensions.isChecked() else None,
            None if self.noSizeLimit.isChecked() else int(self.sizeInput.text()),
            self.formatCombo.currentText(),
            self.keepDimensions.isChecked(),
            self.output_dir_path  # 请确保这是正确的参数
        )

        # 创建WorkerSignals实例
        self.signals = WorkerSignals()
        self.signals.finished.connect(self.updateLabel)  # 连接信号

        # 获取所有图像文件路径
        file_paths = glob.glob(self.source_dir_path + '/**/*.[jp][pn]g', recursive=True)

        
        self.tasksCompleted = 0  # 重置计数器
        self.totalTasks = len(file_paths)  # 设置总任务数量

        for file_path in file_paths:
            rel_path = os.path.relpath(file_path, self.source_dir_path)
            output_file_path = os.path.join(self.output_dir_path, rel_path)
            os.makedirs(os.path.dirname(output_file_path), exist_ok=True)

            # 创建并启动压缩任务
            task = CompressTask(file_path, output_file_path, options, self.signals.finished)
            self.threadpool.start(task)

    def updateLabel(self, message):
        # 更新UI，显示当前处理的文件
        self.label.setText(f"Compressed: {message}")
        self.tasksCompleted += 1  # 更新完成的任务数量
        if self.tasksCompleted == self.totalTasks:
            self.label.setText('所有图片压缩完成')
            self.disableWidgets(False)  # 重新启用控件

    def processImage(self, input_path, output_path, options):
        compress_image(input_path, output_path, options)
        self.updateUI("Compressed: " + input_path)

    def updateUI(self, message):
        # 确保UI更新在主线程中执行
        if threading.current_thread() is threading.main_thread():
            self.label.setText(message)
        else:
            self.label.setText('压缩完成')
            self.disableWidgets(False)

    def onCompressionFinished(self):
        self.disableWidgets(False)  # 重新启用控件
        self.dir_path = ''  # 清空选择的目录
        self.btnStartCompress.setEnabled(False)  # 禁用开始压缩按钮    

    def disableWidgets(self, disable):
        # 禁用或启用所有控件
        self.formatCombo.setDisabled(disable)
        self.widthInput.setDisabled(disable)
        self.keepDimensions.setDisabled(disable)
        self.sizeInput.setDisabled(disable)
        self.noSizeLimit.setDisabled(disable)
        self.btnStartCompress.setDisabled(disable)
        self.btnSelectSourceDirectory.setDisabled(disable)
        self.btnSelectOutputDirectory.setDisabled(disable)


# 主函数
def main():
    app = QApplication(sys.argv)
    ex = ImageCompressor()
    ex.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()