from collections import Counter, deque
from pathlib import Path
import sys

from PIL import Image


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

    samples = []
    border_points = [
        (0, 0), (w - 1, 0), (0, h - 1), (w - 1, h - 1),
        (w // 2, 0), (w // 2, h - 1), (0, h // 2), (w - 1, h // 2),
        (1, 1), (w - 2, 1), (1, h - 2), (w - 2, h - 2),
    ]
    for x, y in border_points:
        samples.append(px[x, y][:3])

    bg = (
        sum(c[0] for c in samples) // len(samples),
        sum(c[1] for c in samples) // len(samples),
        sum(c[2] for c in samples) // len(samples),
    )
    # For checkerboard backgrounds, we also capture another dominant border color.
    border_all = []
    for x in range(w):
        border_all.append(px[x, 0][:3])
        border_all.append(px[x, h - 1][:3])
    for y in range(h):
        border_all.append(px[0, y][:3])
        border_all.append(px[w - 1, y][:3])
    top_border_colors = Counter(border_all).most_common(2)
    bg2 = top_border_colors[1][0] if len(top_border_colors) > 1 else bg
    threshold = 52.0

    visited = [[False] * w for _ in range(h)]
    q = deque()
    for x in range(w):
        q.append((x, 0))
        q.append((x, h - 1))
    for y in range(h):
        q.append((0, y))
        q.append((w - 1, y))

    def dist(a: tuple[int, int, int], b: tuple[int, int, int]) -> float:
        dr = a[0] - b[0]
        dg = a[1] - b[1]
        db = a[2] - b[2]
        return (dr * dr + dg * dg + db * db) ** 0.5

    def is_grayish(rgb: tuple[int, int, int]) -> bool:
        hi = max(rgb)
        lo = min(rgb)
        return (hi - lo) <= 20

    def close_to_bg(rgb: tuple[int, int, int]) -> bool:
        # Match one of checkerboard tones or similar gray edge tone.
        if dist(rgb, bg) <= threshold or dist(rgb, bg2) <= threshold:
            return True
        return is_grayish(rgb) and 65 <= rgb[0] <= 180 and 65 <= rgb[1] <= 180 and 65 <= rgb[2] <= 180

    changed = 0
    while q:
        x, y = q.popleft()
        if x < 0 or x >= w or y < 0 or y >= h or visited[y][x]:
            continue
        visited[y][x] = True
        r, g, b, a = px[x, y]
        if a == 0 or not close_to_bg((r, g, b)):
            continue

        px[x, y] = (r, g, b, 0)
        changed += 1
        q.append((x + 1, y))
        q.append((x - 1, y))
        q.append((x, y + 1))
        q.append((x, y - 1))

    # Second pass: remove remaining checkerboard "islands" anywhere in canvas
    # (not necessarily edge-connected). This targets neutral gray pixels only.
    for y in range(h):
        for x in range(w):
            r, g, b, a = px[x, y]
            if a == 0:
                continue
            mx = max(r, g, b)
            mn = min(r, g, b)
            avg = (r + g + b) // 3
            # Typical checkerboard: neutral gray in mid-dark range.
            if (mx - mn) <= 16 and 35 <= avg <= 185:
                px[x, y] = (r, g, b, 0)
                changed += 1

    img.save(path)
    alpha_hist = Image.open(path).convert("RGBA").getchannel("A").histogram()
    transparent = alpha_hist[0]
    nonzero = sum(alpha_hist[1:])
    return changed, bg, bg2, transparent, nonzero


def main() -> int:
    if len(sys.argv) < 2:
        print("Usage: python make_png_bg_transparent.py <png_path>")
        return 2
    path = Path(sys.argv[1])
    if not path.exists():
        print(f"not found: {path}")
        return 1

    changed, bg, bg2, transparent, nonzero = make_transparent(path)
    print(f"done: {path}")
    print(
        f"bg1={bg}, bg2={bg2}, changed={changed}, "
        f"transparent={transparent}, nonzero_alpha={nonzero}"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())