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

def get_file_kb(file_path: str) -> float:
    """获取文件大小 KB"""
    size_byte = os.path.getsize(file_path)
    return size_byte / 1024

def compress_to_jpg_loop(
    input_path: str,
    output_path: str,
    min_size_kb: float,
    max_size_kb: float,
    init_quality: int,
    quality_step: int,
    min_quality: int,
    max_width: int,
    max_height: int
):
    try:
        ori_size = get_file_kb(input_path)
        print(f"\n原图: {os.path.basename(input_path)} | 大小: {ori_size:.2f} KB")

        # 小于阈值，直接复制原图，不压缩
        if ori_size < min_size_kb:
            shutil.copy2(input_path, output_path)
            print(f"✅ 小于{min_size_kb}KB，无需压缩，直接复制")
            return

        # 打开图片，处理透明通道
        img = Image.open(input_path)
        if img.mode in ("RGBA", "P"):
            new_img = Image.new("RGB", img.size, (255, 255, 255))
            mask = img.split()[-1] if img.mode == "RGBA" else None
            new_img.paste(img, mask=mask)
            img = new_img

        # 等比例缩放
        if max_width and max_height:
            img.thumbnail((max_width, max_height))

        current_quality = init_quality
        while True:
            img.save(output_path, format="JPEG", quality=current_quality, optimize=True)
            out_size = get_file_kb(output_path)

            if out_size <= max_size_kb or current_quality <= min_quality:
                print(f"✅ 最终画质 {current_quality} | 输出大小: {out_size:.2f} KB")
                break

            current_quality -= quality_step
            print(f"⚠️ {out_size:.2f}KB > {max_size_kb}KB，降低画质至 {current_quality} 重新压缩")

    except Exception as e:
        print(f"❌ 处理失败 {input_path}：{str(e)}")


def batch_compress_folder(
    input_dir: str,
    output_dir: str,
    min_kb,
    max_kb,
    init_q,
    q_step,
    min_q,
    max_w,
    max_h
):
    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 not filename.lower().endswith(support_ext):
            continue
        if os.path.isdir(file_path):
            continue

        name_no_ext = os.path.splitext(filename)[0]
        out_file = os.path.join(output_dir, f"{name_no_ext}.jpg")
        compress_to_jpg_loop(
            input_path=file_path,
            output_path=out_file,
            min_size_kb=min_kb,
            max_size_kb=max_kb,
            init_quality=init_q,
            quality_step=q_step,
            min_quality=min_q,
            max_width=max_w,
            max_height=max_h
        )


if __name__ == "__main__":
    # ==================== 批量配置区（仅需修改这里） ====================
    INPUT_FOLDER = "./input"
    OUTPUT_FOLDER = "./output"
    MIN_SIZE_KB = 200      # 小于该KB不压缩
    MAX_SIZE_KB = 800      # 超过该KB循环再压
    INIT_QUALITY = 75      # 初始画质
    QUALITY_STEP = 8       # 每次降低画质数值
    MIN_QUALITY = 20       # 最低画质底线
    LIMIT_W = 2560         # 不需要缩放填 None
    LIMIT_H = 1440
    # ==================================================================

    # 仅批量执行，无任何单张代码
    batch_compress_folder(
        input_dir=INPUT_FOLDER,
        output_dir=OUTPUT_FOLDER,
        min_kb=MIN_SIZE_KB,
        max_kb=MAX_SIZE_KB,
        init_q=INIT_QUALITY,
        q_step=QUALITY_STEP,
        min_q=MIN_QUALITY,
        max_w=LIMIT_W,
        max_h=LIMIT_H
    )