"""
根据抠图结果裁剪多余的透明区域，同时保证整个文件夹内的图片尺寸一致
"""
import sys
import os
from glob import glob
import numpy as np
from PIL import Image

SUPPORTED_FORMATS = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff", ".tif"]

def collect_image_files(folder):
    """收集文件夹中所有支持格式的图片路径"""
    image_paths = []
    for ext in SUPPORTED_FORMATS:
        image_paths.extend(glob(os.path.join(folder, f"*{ext}")))
        image_paths.extend(glob(os.path.join(folder, f"*{ext.upper()}")))
    # 去重并排序（Windows文件系统不区分大小写）
    return sorted(set(image_paths))

def compute_global_bbox(image_paths, alpha_threshold=0):
    """
    计算所有图片中非透明像素的全局 bounding box
    Returns: (left, top, right, bottom)
    """
    min_x = None
    min_y = None
    max_x = None
    max_y = None

    for path in image_paths:
        try:
            image = Image.open(path).convert("RGBA")
            alpha = np.array(image)[:, :, 3]
            mask = alpha > alpha_threshold
            if not np.any(mask):
                continue

            rows, cols = np.where(mask)
            left, right = int(cols.min()), int(cols.max())
            top, bottom = int(rows.min()), int(rows.max())

            min_x = left if min_x is None else min(min_x, left)
            min_y = top if min_y is None else min(min_y, top)
            max_x = right if max_x is None else max(max_x, right)
            max_y = bottom if max_y is None else max(max_y, bottom)
        except Exception as error:
            print(f"读取 {path} 时出错: {error}")

    if min_x is None or min_y is None or max_x is None or max_y is None:
        return None

    # PIL 的 crop 右下角是开区间，因此 +1
    return (min_x, min_y, max_x + 1, max_y + 1)

def crop_images(input_folder, output_folder):
    """裁剪输入文件夹中的图片并输出到指定文件夹"""
    print('\n' + '=' * 60)
    print('[Step 2/2] Smart Cropping')
    print('=' * 60)
    
    if not os.path.exists(input_folder):
        print(f"[X] Error: Input folder not found: {input_folder}")
        return 0

    print(f"-> Input folder: {input_folder}")
    
    image_paths = collect_image_files(input_folder)
    if not image_paths:
        print(f"[X] Error: No supported image files found")
        return 0

    print(f"[OK] Found {len(image_paths)} images")
    print(f"-> Computing global crop area...")

    bbox = compute_global_bbox(image_paths)
    if bbox is None:
        print("[X] Warning: No valid non-transparent pixels detected")
        return 0

    left, top, right, bottom = bbox
    target_width = right - left
    target_height = bottom - top

    print(f"[OK] Crop area: left={left}, top={top}, width={target_width}, height={target_height}")

    os.makedirs(output_folder, exist_ok=True)
    print(f"-> Output folder: {output_folder}")
    print('-' * 60)

    success = 0
    for idx, path in enumerate(image_paths, 1):
        try:
            filename = os.path.basename(path)
            print(f"[{idx}/{len(image_paths)}] Cropping: {filename}")
            print(f'PROGRESS: {idx}/{len(image_paths)}')
            
            image = Image.open(path).convert("RGBA")
            cropped = image.crop(bbox)

            # 确保裁剪后的尺寸一致
            if cropped.size != (target_width, target_height):
                fixed = Image.new("RGBA", (target_width, target_height), (0, 0, 0, 0))
                fixed.paste(cropped, (0, 0), cropped)
                cropped = fixed

            output_path = os.path.join(output_folder, filename)
            cropped.save(output_path, "PNG")
            print(f"    [OK] Done")
            success += 1
        except Exception as error:
            print(f"    [X] Failed: {error}")

    print('\n' + '=' * 60)
    print(f"[Step 2 Complete] Success: {success}/{len(image_paths)}")
    print('=' * 60)
    return success

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("用法: python cut-mini-size.py <输入文件夹> <输出文件夹>")
        sys.exit(1)
    
    input_folder = sys.argv[1]
    output_folder = sys.argv[2]
    
    processed = crop_images(input_folder, output_folder)
    sys.exit(0 if processed > 0 else 1)