#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""AI 插画后处理：
① 去除服务端强制水印（右下角「AI生成」灰胶囊）—— 按亮度/饱和度定位后以背景底色回填
② 长边降采样到 1280（控体积）
③ 输出 webp（内嵌用）+ clean png（存档）
原始 ai_<name>.png 保留不动。
"""
import io
import os
import sys

from PIL import Image, ImageDraw, ImageFilter

OUT = '/Users/liuxinyuan/Desktop/Hermes输出-工作类/培训与上手/培训/_img_price_deck'
NAMES = ['scene_tug', 'scene_weigh', 'scene_deal', 'scene_price', 'jobs_cart']
LONG = 1280


def lum_sat(p):
    r, g, b = p[0], p[1], p[2]
    return (0.299 * r + 0.587 * g + 0.114 * b), (max(r, g, b) - min(r, g, b))


def find_watermark(im):
    """只在紧贴右下角的小窗口内找灰胶囊，避免把地面阴影/浅灰形状误判为水印。"""
    W, H = im.size
    x0, y0 = int(W * 0.865), int(H * 0.855)
    px = im.load()
    xs, ys = [], []
    for y in range(y0, H):
        for x in range(x0, W):
            L, S = lum_sat(px[x, y])
            if S < 30 and 110 < L < 218:            # 灰胶囊（浅灰、低饱和）
                xs.append(x)
                ys.append(y)
    if not xs:
        return None
    box = (min(xs), min(ys), max(xs), max(ys))
    # 尺寸异常大 → 说明误判，退回固定角位
    if (box[2] - box[0]) > 0.17 * W or (box[3] - box[1]) > 0.15 * H:
        box = (int(W * 0.875), int(H * 0.865), int(W * 0.99), int(H * 0.965))
    return box


def clean(name):
    src = os.path.join(OUT, 'ai_%s.png' % name)
    im = Image.open(src).convert('RGB')
    W, H = im.size
    box = find_watermark(im)
    print('  %-12s %dx%d  watermark bbox=%s' % (name, W, H, box))
    if box:
        bx0, by0, bx1, by1 = box
        pad = 10
        bx0, by0 = max(0, bx0 - pad), max(0, by0 - pad)
        bx1, by1 = min(W, bx1 + pad), min(H, by1 + pad)
        # 底色取「同高、左侧」一段纯背景的中位色，保留原图色调
        ref = im.crop((max(0, bx0 - 260), by0, max(1, bx0 - 40), by1))
        cols = sorted(r for r, g, b in ref.getdata())
        pix = list(ref.getdata())
        med = tuple(sorted(c[i] for c in pix)[len(pix) // 2] for i in range(3))
        # 回填 + 轻噪声，避免出现纯色方块边缘
        patch = Image.new('RGB', (bx1 - bx0, by1 - by0), med)
        d = ImageDraw.Draw(patch)
        pr = (bx1 - bx0) // 14
        for i in range(0, bx1 - bx0, max(3, pr)):
            for j in range(0, by1 - by0, max(3, pr)):
                v = ((i * 7 + j * 13) % 5) - 2
                d.point((i, j), fill=(med[0] + v, med[1] + v, med[2] + v))
        patch = patch.filter(ImageFilter.GaussianBlur(1.1))
        im.paste(patch, (bx0, by0))
    # 降采样
    w, h = im.size
    k = LONG / float(max(w, h))
    if k < 1:
        im = im.resize((max(1, int(round(w * k))), max(1, int(round(h * k)))),
                       Image.LANCZOS)
    cpng = os.path.join(OUT, 'ai_%s_clean.png' % name)
    im.save(cpng)
    webp = os.path.join(OUT, 'ai_%s_clean.webp' % name)
    im.save(webp, 'WEBP', quality=86, method=6)
    print('      clean %dx%d  png %dKB  webp %dKB' % (
        im.size[0], im.size[1], os.path.getsize(cpng) // 1024,
        os.path.getsize(webp) // 1024))
    return im.size


if __name__ == '__main__':
    print('后处理 →', OUT)
    for n in NAMES:
        clean(n)