"""Remove white background from an image by flood-filling from edges only."""

import sys
from collections import deque
from PIL import Image

def remove_white_bg(path: str, tolerance: int = 20, out: str | None = None):
    img = Image.open(path).convert("RGBA")
    pixels = img.load()
    w, h = img.size
    threshold = 255 - tolerance

    visited = set()
    queue = deque()

    # Seed from all edge pixels
    for x in range(w):
        queue.append((x, 0))
        queue.append((x, h - 1))
    for y in range(h):
        queue.append((0, y))
        queue.append((w - 1, y))

    while queue:
        x, y = queue.popleft()
        if (x, y) in visited or x < 0 or y < 0 or x >= w or y >= h:
            continue
        visited.add((x, y))
        r, g, b, a = pixels[x, y]
        if r >= threshold and g >= threshold and b >= threshold:
            pixels[x, y] = (r, g, b, 0)
            queue.append((x + 1, y))
            queue.append((x - 1, y))
            queue.append((x, y + 1))
            queue.append((x, y - 1))

    # Crop to bounding box of non-transparent pixels
    bbox = img.getbbox()
    if bbox:
        img = img.crop(bbox)

    dest = out or path
    img.save(dest)
    print(f"Saved to {dest} ({img.size[0]}x{img.size[1]})")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python remove_bg.py <image> [tolerance] [output]")
        sys.exit(1)
    path = sys.argv[1]
    tol = int(sys.argv[2]) if len(sys.argv) > 2 else 20
    out = sys.argv[3] if len(sys.argv) > 3 else None
    remove_white_bg(path, tol, out)