from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image, ImageFilter


def clean_alpha(src: Path, dst: Path) -> tuple[int, int]:
    img = Image.open(src).convert("RGBA")
    w, h = img.size
    px = img.load()

    alpha = img.getchannel("A")
    solid = alpha.point(lambda a: 255 if a > 8 else 0)
    inner = solid.filter(ImageFilter.MinFilter(21))
    edge = Image.new("L", img.size, 0)
    edge.paste(solid)
    edge_px = edge.load()
    inner_px = inner.load()

    removed = 0
    softened = 0
    for y in range(h):
        upper_zone = y < int(h * 0.55)
        for x in range(w):
            r, g, b, a = px[x, y]
            if a == 0:
                continue

            near_edge = edge_px[x, y] and not inner_px[x, y]
            if not near_edge and a > 225:
                continue

            bright = (r + g + b) / 3
            chroma = max(r, g, b) - min(r, g, b)
            semi = a < 245

            neutral_matte = semi and bright > 118 and chroma < 82
            pink_or_cyan_matte = (
                semi
                and upper_zone
                and bright > 130
                and (
                    (r > 135 and b > 130 and g > 80)
                    or (g > 135 and b > 130 and r < 175)
                )
                and chroma < 135
            )

            if near_edge and (neutral_matte or pink_or_cyan_matte):
                px[x, y] = (r, g, b, 0)
                removed += 1
            elif semi and neutral_matte:
                px[x, y] = (r, g, b, int(a * 0.2))
                softened += 1

    dst.parent.mkdir(parents=True, exist_ok=True)
    img.save(dst)
    return removed, softened


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("paths", nargs="+", type=Path)
    parser.add_argument("--in-place", action="store_true")
    args = parser.parse_args()

    for path in args.paths:
        dst = path if args.in_place else path.with_name(path.stem + "_clean.png")
        removed, softened = clean_alpha(path, dst)
        print(f"{path} -> {dst} removed={removed} softened={softened}")


if __name__ == "__main__":
    main()