"""
使用 rembg 库进行图像背景移除
"""
import sys
import os
from rembg import remove
from PIL import Image
from glob import glob

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

def extract_character_rembg(input_path, output_path):
    """
    使用 rembg 库移除图像背景
    """
    try:
        print(f'  -> 读取图片: {os.path.basename(input_path)}')
        
        if not os.path.exists(input_path):
            print(f'    [X] 错误: 文件不存在')
            return False
        
        # 创建输出目录
        output_dir = os.path.dirname(output_path)
        if output_dir and not os.path.exists(output_dir):
            print(f'  -> 创建输出目录: {output_dir}')
            os.makedirs(output_dir, exist_ok=True)
        
        # 读取输入图像
        print(f'  -> 加载图像...')
        input_image = Image.open(input_path)
        original_size = input_image.size
        print(f'    图像尺寸: {original_size[0]}x{original_size[1]}')
        
        # 确保输入图像是RGB模式
        if input_image.mode != 'RGB':
            print(f'  -> 转换颜色模式: {input_image.mode} -> RGB')
            input_image = input_image.convert('RGB')
        
        # 使用 rembg 移除背景
        print(f'  -> 执行AI抠图...')
        output_image = remove(input_image, post_process_mask=False)
        print(f'    [OK] 抠图完成')
        
        # 确保输出图像尺寸与输入图像一致
        if output_image.size != original_size:
            print(f'  -> 调整图像尺寸')
            final_image = Image.new('RGBA', original_size, (0, 0, 0, 0))
            final_image.paste(output_image, (0, 0), output_image if output_image.mode == 'RGBA' else None)
            output_image = final_image
        else:
            if output_image.mode != 'RGBA':
                output_image = output_image.convert('RGBA')
        
        # 保存结果
        print(f'  -> 保存结果: {os.path.basename(output_path)}')
        output_image.save(output_path, 'PNG')
        print(f'    [OK] 保存成功')
        return True
        
    except Exception as error:
        print(f'    [X] 处理失败: {error}')
        return False

def process_folder(input_folder, output_folder):
    """
    批量处理文件夹中的所有图片
    """
    print('=' * 60)
    print('[Step 1/2] AI Image Matting')
    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}')
    
    # 创建输出文件夹
    if not os.path.exists(output_folder):
        print(f'-> Creating output folder: {output_folder}')
        os.makedirs(output_folder, exist_ok=True)
    else:
        print(f'-> Output folder: {output_folder}')
    
    # 获取所有支持的图片文件
    print(f'-> Scanning image files...')
    image_files = []
    for ext in SUPPORTED_FORMATS:
        image_files.extend(glob(os.path.join(input_folder, f'*{ext}')))
        image_files.extend(glob(os.path.join(input_folder, f'*{ext.upper()}')))
    
    # 去重（Windows文件系统不区分大小写，可能重复）
    image_files = list(set(image_files))
    image_files.sort()
    
    if not image_files:
        print(f'[X] No supported image files found')
        return 0
    
    print(f'[OK] Found {len(image_files)} images')
    print('-' * 60)
    
    success_count = 0
    for idx, input_path in enumerate(image_files, 1):
        filename = os.path.basename(input_path)
        name_without_ext = os.path.splitext(filename)[0]
        output_path = os.path.join(output_folder, f'{name_without_ext}.png')
        
        print(f'\n[{idx}/{len(image_files)}] {filename}')
        print(f'PROGRESS: {idx}/{len(image_files)}')
        
        if extract_character_rembg(input_path, output_path):
            success_count += 1
            print(f'    [OK] Done')
    
    print('\n' + '=' * 60)
    print(f'[Step 1 Complete] Success: {success_count}/{len(image_files)}')
    print('=' * 60)
    return success_count

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