import os
from PIL import Image, ImageDraw, ImageFont

# --- 配置区域 ---
TARGET_FOLDER = r"D:\1-working\1-tagging\1-new-round\Export 1DEU-Round2-800-wasser\upload-compare"
WATERMARK_TEXT = "myPaintLab"
TARGET_SIZE = (800, 800)
# ----------------

def process_images_inplace():
    if not os.path.exists(TARGET_FOLDER):
        print(f"❌ 错误：文件夹不存在 - {TARGET_FOLDER}")
        return

    # 获取所有图片文件
    all_files = [f for f in os.listdir(TARGET_FOLDER) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
    total = len(all_files)
    
    print(f"🚀 开始处理：文件夹内总计 {total} 张图片（统一白色背景）...")

    for index, filename in enumerate(all_files):
        file_path = os.path.join(TARGET_FOLDER, filename)
        
        try:
            with Image.open(file_path) as img:
                # 记录原图格式 (如 'PNG', 'JPEG')
                original_format = img.format
                
                # --- 1. 处理原图透明度 ---
                # 为了统一白色背景，先转为 RGBA 处理
                img = img.convert("RGBA")
                
                # --- 2. 缩放逻辑 ---
                width, height = img.size
                scale = 800 / max(width, height)
                new_w, new_h = int(width * scale), int(height * scale)
                resized_img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
                
                # --- 3. 创建纯白色画布 ---
                # 无论原图是什么，底色一律白色
                white_canvas = Image.new("RGBA", TARGET_SIZE, (255, 255, 255, 255))
                
                # 居中粘贴 (resized_img 作为 mask 可以正确处理原图的透明部分，使其显示白色底色)
                paste_pos = ((TARGET_SIZE[0] - new_w) // 2, (TARGET_SIZE[1] - new_h) // 2)
                white_canvas.paste(resized_img, paste_pos, mask=resized_img)

                # --- 4. 水印逻辑 ---
                overlay = Image.new("RGBA", TARGET_SIZE, (255, 255, 255, 0))
                draw = ImageDraw.Draw(overlay)
                
                try:
                    font = ImageFont.truetype("arial.ttf", 45)
                except:
                    font = ImageFont.load_default()
                
                bbox = draw.textbbox((0, 0), WATERMARK_TEXT, font=font)
                tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
                text_position = ((800 - tw) // 2, (800 - th) // 2)
                draw.text(text_position, WATERMARK_TEXT, font=font, fill=(150, 150, 150, 45))
                
                # 合成水印
                final_img = Image.alpha_composite(white_canvas, overlay)
                
                # --- 5. 保存处理（核心修复） ---
                # 如果原格式是 JPEG/JPG，必须转回 RGB 模式（不能带 A 通道），否则保存会报错
                # 如果原格式是 PNG，保持 RGBA 或转 RGB 均可（这里统一转 RGB 除非你需要保留 PNG 但不需要透明）
                if original_format in ['JPEG', 'JPG']:
                    final_img = final_img.convert("RGB")
                else:
                    # 如果你希望 PNG 也完全没有透明通道，统一转 RGB
                    final_img = final_img.convert("RGB")

                # 按照记录的原格式保存
                final_img.save(file_path, format=original_format)

            print(f"[{index+1}/{total}] 已处理: {filename}")

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

    print("\n✅ 所有图片已处理：800x800 像素，统一白色背景，原格式覆盖。")

if __name__ == "__main__":
    process_images_inplace()