# pip install pillow
from PIL import Image
import os
import shutil

SUPPORTED_OUTPUT_FORMATS = ("jpg", "jpeg", "png", "webp")


def get_file_kb(file_path: str) -> float:
    return os.path.getsize(file_path) / 1024


def _save_image(img: Image.Image, output_path: str, output_format: str, quality: int):
    if output_format == "png":
        img.save(output_path, format="PNG", optimize=True, compress_level=9)
    elif output_format == "webp":
        img.save(output_path, format="WEBP", quality=quality, method=6)
    else:
        img.save(output_path, format="JPEG", quality=quality, optimize=True)


def compress_to_format_loop(input_path, output_path, min_size_kb, max_size_kb,
                            init_quality, quality_step, min_quality,
                            max_width, max_height, output_format):
    try:
        # None 表示保持输入图片格式。
        if output_format is None:
            output_format = os.path.splitext(input_path)[1].lower().lstrip(".")
        output_format = output_format.lower().lstrip(".")
        if output_format not in SUPPORTED_OUTPUT_FORMATS:
            raise ValueError(f"不支持的输出格式: {output_format}")

        ori_size = get_file_kb(input_path)
        print(f"\n原图: {os.path.basename(input_path)} | 大小: {ori_size:.2f} KB")

        # INIT_QUALITY=None：不压缩、不缩放；目标格式不同则只做格式转换。
        # 小于阈值：同样不压缩、不缩放；目标格式不同则只做格式转换。
        if init_quality is None or ori_size < min_size_kb:
            if os.path.splitext(input_path)[1].lower().lstrip(".") == output_format:
                shutil.copy2(input_path, output_path)
                print(f"✅ 小于{min_size_kb}KB，无需压缩，直接复制")
            else:
                img = Image.open(input_path)
                img = _prepare_image(img, output_format)
                _save_image(img, output_path, output_format, init_quality)
                reason = "INIT_QUALITY=None" if init_quality is None else f"小于{min_size_kb}KB"
                print(f"✅ {reason}，不压缩，仅转换为 {output_format}")
            return

        img = Image.open(input_path)
        img = _prepare_image(img, output_format)
        if max_width and max_height:
            img.thumbnail((max_width, max_height))

        # PNG 没有 quality 参数，直接用最高无损压缩保存。
        if output_format == "png":
            _save_image(img, output_path, output_format, init_quality)
            print(f"✅ PNG 已保存 | 输出大小: {get_file_kb(output_path):.2f} KB")
            return

        quality = init_quality
        while True:
            _save_image(img, output_path, output_format, quality)
            out_size = get_file_kb(output_path)
            if out_size <= max_size_kb or quality <= min_quality:
                print(f"✅ 最终画质 {quality} | 输出大小: {out_size:.2f} KB")
                break
            quality = max(min_quality, quality - quality_step)
            print(f"⚠️ {out_size:.2f}KB > {max_size_kb}KB，降低画质至 {quality} 重新压缩")
    except Exception as e:
        print(f"❌ 处理失败 {input_path}：{str(e)}")


def _prepare_image(img, output_format):
    has_alpha = "A" in img.getbands() or (img.mode == "P" and "transparency" in img.info)
    if output_format in ("jpg", "jpeg"):
        if has_alpha:
            rgba = img.convert("RGBA")
            result = Image.new("RGB", rgba.size, (255, 255, 255))
            result.paste(rgba, mask=rgba.getchannel("A"))
            return result
        return img.convert("RGB")
    return img.convert("RGBA" if has_alpha else "RGB")


def batch_compress_folder(input_dir, output_dir, min_kb, max_kb, init_q,
                          q_step, min_q, max_w, max_h, output_format):
    os.makedirs(output_dir, exist_ok=True)
    support_ext = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
    for filename in os.listdir(input_dir):
        file_path = os.path.join(input_dir, filename)
        if os.path.isdir(file_path) or not filename.lower().endswith(support_ext):
            continue
        name_no_ext = os.path.splitext(filename)[0]
        suffix = output_format or os.path.splitext(filename)[1].lstrip(".")
        out_file = os.path.join(output_dir, f"{name_no_ext}.{suffix}")
        compress_to_format_loop(file_path, out_file, min_kb, max_kb, init_q,
                                q_step, min_q, max_w, max_h, output_format)


if __name__ == "__main__":
    # 输入输出文件夹
    INPUT_FOLDER = "./input"
    OUTPUT_FOLDER = "./output"
    # 首次压缩质量，不压缩填 None
    INIT_QUALITY = 75
    # 过小不压缩，过大重复压缩
    MIN_SIZE_KB = 100
    MAX_SIZE_KB = 500
    # 重复压缩质量递减及最小值
    QUALITY_STEP = 8
    MIN_QUALITY = 20
    # 更换格式，不转换格式填 None
    OUTPUT_FORMAT = "webp"
    # 尺寸压缩，不缩放填 None
    LIMIT_W = 2560
    LIMIT_H = 1440

    batch_compress_folder(INPUT_FOLDER, OUTPUT_FOLDER, MIN_SIZE_KB, MAX_SIZE_KB,
                          INIT_QUALITY, QUALITY_STEP, MIN_QUALITY, LIMIT_W,
                          LIMIT_H, OUTPUT_FORMAT)