import os
import random
from pathlib import Path
from PIL import Image

# 配置参数
MAIN_BG_DIR = './微信无头像主界面截图'
AVATAR_DIR = './微信头像素材'
OUTPUT_DIR = './output'

# 区域参数
REGION_X = 8
REGION_Y = 31
REGION_W = 75
REGION_H = 75
AVATAR_SIZE = 54  # 54x54 px

# 输出目录准备
def ensure_output_dir():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

def get_image_files(folder):
    exts = ['.jpg', '.jpeg', '.png']
    return [f for f in Path(folder).iterdir() if f.suffix.lower() in exts]

def process_avatar(avatar_path):
    """缩放/裁剪为54x54，保持居中"""
    img = Image.open(avatar_path).convert('RGBA')
    # 等比缩放
    min_side = min(img.width, img.height)
    # 居中裁剪为正方形
    left = (img.width - min_side) // 2
    top = (img.height - min_side) // 2
    img = img.crop((left, top, left + min_side, top + min_side))
    img = img.resize((AVATAR_SIZE, AVATAR_SIZE), Image.LANCZOS)
    return img

def random_offset():
    x = random.randint(REGION_X, REGION_X + REGION_W - AVATAR_SIZE)
    y = random.randint(REGION_Y, REGION_Y + REGION_H - AVATAR_SIZE)
    return x, y

def main():
    ensure_output_dir()
    main_imgs = get_image_files(MAIN_BG_DIR)
    avatar_imgs = get_image_files(AVATAR_DIR)
    if not main_imgs or not avatar_imgs:
        print('未找到主界面或头像素材图片')
        return
    avatar_count = 0
    for avatar_idx, avatar_path in enumerate(avatar_imgs, 1):
        avatar_img = process_avatar(avatar_path)
        avatar_name = avatar_path.stem
        for main_idx, main_path in enumerate(main_imgs, 1):
            with Image.open(main_path).convert('RGBA') as main_img:
                x, y = random_offset()
                composed = main_img.copy()
                composed.paste(avatar_img, (x, y), avatar_img)
                # 命名规则
                out_name = f"avatar_{avatar_idx:02d}_{x}-{y}_{AVATAR_SIZE}x{AVATAR_SIZE}.jpg"
                out_path = os.path.join(OUTPUT_DIR, out_name)
                # 转为RGB再保存为jpg
                composed.convert('RGB').save(out_path, 'JPEG', quality=95)
                print(f"保存: {out_path}")
                avatar_count += 1
    print(f"共生成{avatar_count}张合成图片。")

if __name__ == '__main__':
    main()