#!/usr/bin/env python3
"""
图片处理脚本
裁剪、缩放、合成背景

用法：
  python3 process_image.py input.jpg output.jpg --ratio 1:1 --style white
  python3 process_image.py input.jpg output.jpg --ratio 4:3 --style minimal
"""

import argparse
import json
import os
import sys

from PIL import Image, ImageFilter, ImageDraw

# 目标尺寸（长边）
TARGET_SIZE = 800

# 背景色
BG_COLORS = {
    "white": (255, 255, 255),
    "light_gray": (245, 245, 245),
    "transparent": None,
}

# 比例映射
RATIO_MAP = {
    "1:1": (1, 1),
    "4:3": (4, 3),
    "3:4": (3, 4),
    "16:9": (16, 9),
}

# 水印检测阈值：像素亮度 > 此值且 < 255 的区域可能是水印
WATERMARK_BRIGHT_MIN = 200
WATERMARK_BRIGHT_MAX = 254
# 边缘裁剪比例（去除角标水印）
EDGE_CROP_RATIO = 0.03


def crop_edges(img, ratio=EDGE_CROP_RATIO):
    """裁剪边缘，去除角落水印/标识"""
    w, h = img.size
    cx = int(w * ratio)
    cy = int(h * ratio)
    return img.crop((cx, cy, w - cx, h - cy))


def remove_semi_transparent_watermark(img):
    """去除半透明文字水印（常见于素材网站）"""
    if img.mode == "RGB":
        img = img.convert("RGBA")

    pixels = img.load()
    w, h = img.size

    # 扫描检测亮度异常高的像素（接近白色但不是纯白 → 可能是半透明白色水印）
    for y in range(h):
        for x in range(w):
            r, g, b, a = pixels[x, y]
            # 计算亮度
            brightness = (r + g + b) / 3
            # 如果亮度很高但不是纯白，且 RGB 三通道接近（灰白色），可能是水印
            if WATERMARK_BRIGHT_MIN <= brightness <= WATERMARK_BRIGHT_MAX:
                if abs(r - g) < 15 and abs(g - b) < 15 and abs(r - b) < 15:
                    # 取周围像素的平均值来替换
                    pixels[x, y] = _get_neighbor_avg(pixels, x, y, w, h)

    return img.convert("RGB")


def _get_neighbor_avg(pixels, x, y, w, h, radius=3):
    """获取周围非水印像素的平均颜色"""
    total_r, total_g, total_b, count = 0, 0, 0, 0

    for dy in range(-radius, radius + 1):
        for dx in range(-radius, radius + 1):
            nx, ny = x + dx, y + dy
            if 0 <= nx < w and 0 <= ny < h and (dx != 0 or dy != 0):
                r, g, b, a = pixels[nx, ny]
                brightness = (r + g + b) / 3
                # 只统计不是水印的像素
                if brightness < WATERMARK_BRIGHT_MIN or brightness > WATERMARK_BRIGHT_MAX:
                    total_r += r
                    total_g += g
                    total_b += b
                    count += 1

    if count > 0:
        return (total_r // count, total_g // count, total_b // count, 255)
    # 如果周围都是水印，用白色填充
    return (255, 255, 255, 255)


def remove_watermark(img):
    """综合水印去除：边缘裁剪 + 半透明水印清理"""
    # 1. 裁剪边缘（去除角落 logo/水印）
    img = crop_edges(img)

    # 2. 去除半透明文字水印
    img = remove_semi_transparent_watermark(img)

    return img


def crop_to_ratio(img, ratio_str):
    """按比例居中裁剪"""
    w, h = img.size
    rw, rh = RATIO_MAP.get(ratio_str, (1, 1))
    target_ratio = rw / rh
    current_ratio = w / h

    if current_ratio > target_ratio:
        # 图片太宽，裁剪左右
        new_w = int(h * target_ratio)
        left = (w - new_w) // 2
        img = img.crop((left, 0, left + new_w, h))
    elif current_ratio < target_ratio:
        # 图片太高，裁剪上下
        new_h = int(w / target_ratio)
        top = (h - new_h) // 2
        img = img.crop((0, top, w, top + new_h))

    return img


def resize_to_target(img, target_size=TARGET_SIZE):
    """缩放到目标尺寸（长边为 target_size）"""
    w, h = img.size
    if w >= h:
        ratio = target_size / w
    else:
        ratio = target_size / h

    new_w = int(w * ratio)
    new_h = int(h * ratio)
    return img.resize((new_w, new_h), Image.LANCZOS)


def add_white_bg(img):
    """添加白色背景（处理透明通道）"""
    if img.mode in ("RGBA", "LA", "PA"):
        bg = Image.new("RGB", img.size, (255, 255, 255))
        if img.mode == "RGBA":
            bg.paste(img, mask=img.split()[3])
        else:
            bg.paste(img)
        return bg
    elif img.mode != "RGB":
        return img.convert("RGB")
    return img


def add_colored_bg(img, color):
    """添加纯色背景"""
    bg = Image.new("RGB", img.size, color)
    if img.mode in ("RGBA", "LA", "PA"):
        if img.mode == "RGBA":
            bg.paste(img, mask=img.split()[3])
        else:
            bg.paste(img)
    else:
        bg.paste(img)
    return bg


def sharpen(img):
    """轻微锐化"""
    return img.filter(ImageFilter.SHARPEN)


def process_image(input_path, output_path, ratio="1:1", style="white"):
    """完整处理流程"""
    try:
        img = Image.open(input_path)

        # 1. 转为 RGB（如果是 CMYK 等）
        if img.mode not in ("RGB", "RGBA", "LA", "PA"):
            img = img.convert("RGB")

        # 2. 去水印
        img = remove_watermark(img)

        # 3. 按比例裁剪
        img = crop_to_ratio(img, ratio)

        # 4. 缩放
        img = resize_to_target(img, TARGET_SIZE)

        # 5. 合成背景
        if style == "white":
            img = add_white_bg(img)
        elif style == "minimal":
            img = add_white_bg(img)  # minimal 也用白底
        elif style == "scene":
            # 场景图保留原背景，只做裁剪缩放
            if img.mode != "RGB":
                img = add_white_bg(img)
        elif style == "realistic":
            if img.mode != "RGB":
                img = add_white_bg(img)
        else:
            img = add_white_bg(img)

        # 6. 锐化
        img = sharpen(img)

        # 7. 保存
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        img.save(output_path, "JPEG", quality=90)

        return output_path
    except Exception as e:
        print(json.dumps({"error": str(e)}), file=sys.stderr)
        return None


def main():
    parser = argparse.ArgumentParser(description="图片处理")
    parser.add_argument("input", help="输入图片路径")
    parser.add_argument("output", help="输出图片路径")
    parser.add_argument("--ratio", default="1:1", choices=RATIO_MAP.keys(), help="目标比例")
    parser.add_argument("--style", default="white", choices=BG_COLORS.keys(), help="背景风格")
    args = parser.parse_args()

    result = process_image(args.input, args.output, args.ratio, args.style)

    output = {
        "input": args.input,
        "output": result,
        "ratio": args.ratio,
        "style": args.style,
        "success": result is not None,
    }
    print(json.dumps(output, ensure_ascii=False))


if __name__ == "__main__":
    main()