"""dechecker.py —— 兜底用：把 AI 图里画死的"透明棋盘格/背景"抠掉，恢复真透明。

重要：现在生图管线已直接向接口请求真·透明 PNG（providers + pipeline），
正常情况下根本不需要本模块。只有当上游模型仍把背景画成棋盘格时，才作为
最后兜底自动运行。

关键改进（避免戳洞）：只移除"与图像边缘相连"的背景（四周的棋盘格 + 透明边），
角色内部的白色高光/亮点因为被角色包住、不与边缘相连，会被完整保留——
不会再像以前那样把白色切成缺口。
"""
import os
import sys
import numpy as np
from PIL import Image, ImageFilter


def _border_outside(bg):
    """返回与图像四边相连的背景区域(布尔)。优先用 cv2，没有则用 numpy 形态学重建。"""
    try:
        import cv2
        n, labels = cv2.connectedComponents(bg.astype(np.uint8), connectivity=4)
        border = set(labels[0, :]).union(labels[-1, :], labels[:, 0], labels[:, -1])
        border.discard(0)
        return np.isin(labels, list(border))
    except Exception:
        # numpy 兜底：从四边种子在 bg 内反复膨胀直到稳定（形态学重建）
        seed = np.zeros_like(bg)
        seed[0, :] = bg[0, :]; seed[-1, :] = bg[-1, :]
        seed[:, 0] = bg[:, 0]; seed[:, -1] = bg[:, -1]
        for _ in range(4000):
            grown = seed | np.roll(seed, 1, 0) | np.roll(seed, -1, 0) \
                         | np.roll(seed, 1, 1) | np.roll(seed, -1, 1)
            grown &= bg
            if np.array_equal(grown, seed):
                break
            seed = grown
        return seed


def clean(im):
    """返回 (新图RGBA, 抹掉的像素数)。无棋盘/已透明的图基本无操作。"""
    img = im.convert("RGBA")
    arr = np.asarray(img).copy()
    al = arr[:, :, 3]
    rgb = arr[:, :, :3].astype(np.int16)
    bright = rgb.max(2)
    sat = bright - rgb.min(2)
    light = (al >= 200) & (bright > 185) & (sat < 28)   # 棋盘格的白/浅灰方块
    trans = al < 40
    if light.mean() < 0.04:                              # 几乎没有浅色块 → 不是棋盘
        return img, 0
    bg = light | trans
    outside = _border_outside(bg)
    # 只在"边缘相连的浅色块"占比可观时才动手（避免误伤纯色背景插画）
    if (outside & light).mean() < 0.03:
        return img, 0
    # 向内吃 2px 去掉灰边
    m = outside
    for _ in range(2):
        m = (m | np.roll(m, 1, 0) | np.roll(m, -1, 0)
               | np.roll(m, 1, 1) | np.roll(m, -1, 1))
    removed = int(((al > 10) & m).sum())
    a = al.copy(); a[m] = 0
    arr[:, :, 3] = a
    arr[m, 0] = 0; arr[m, 1] = 0; arr[m, 2] = 0
    out = Image.fromarray(arr.astype(np.uint8), "RGBA")
    out.putalpha(out.getchannel("A").filter(ImageFilter.GaussianBlur(1.0)))   # 柔边
    return out, removed


def declutter_img(img, feather=False):
    out, removed = clean(img)
    return out.convert("RGBA"), removed > 0


def declutter_file(src, dst=None):
    out, removed = clean(Image.open(src))
    if removed > 0:
        out.save(dst or src)
    return removed > 0


def _walk(paths):
    for p in paths:
        if os.path.isdir(p):
            for root, _d, files in os.walk(p):
                for f in files:
                    if f.lower().endswith(".png"):
                        yield os.path.join(root, f)
        elif p.lower().endswith(".png"):
            yield p


if __name__ == "__main__":
    n = 0
    for f in _walk(sys.argv[1:] or ["."]):
        try:
            if declutter_file(f):
                n += 1
                print("cleaned:", f)
        except Exception as e:
            print("skip", f, e)
    print(f"done. {n} file(s) cleaned.")