#!/usr/bin/env python3
"""Remove green screen (#00FF00) from an image and produce a transparent PNG.

Usage:
    python3 remove_greenscreen.py input.png output.png [--tolerance 60] [--edge-blur 1]

Dependencies: Pillow (pip3 install Pillow)
"""
import argparse
import sys
from pathlib import Path

try:
    from PIL import Image, ImageFilter
except ImportError:
    print("ERROR: Pillow is required. Install with: pip3 install Pillow", file=sys.stderr)
    sys.exit(1)


def remove_green(input_path: str, output_path: str, tolerance: int = 60, edge_blur: int = 1, fmt: str = "png") -> None:
    img = Image.open(input_path).convert("RGBA")
    pixels = img.load()
    width, height = img.size

    for y in range(height):
        for x in range(width):
            r, g, b, a = pixels[x, y]
            # Check if pixel is "green enough" to be chroma key
            if g > 100 and (g - r) > tolerance and (g - b) > tolerance:
                pixels[x, y] = (0, 0, 0, 0)
            else:
                # De-spill: reduce green tint on edge pixels near transparent areas
                max_rb = max(r, b)
                if g > max_rb + 20:
                    g_corrected = min(g, max_rb + 10)
                    pixels[x, y] = (r, g_corrected, b, a)

    # Edge anti-aliasing via alpha channel blur
    if edge_blur > 0:
        alpha = img.split()[3]
        alpha = alpha.filter(ImageFilter.GaussianBlur(radius=edge_blur))
        img.putalpha(alpha)

    save_format = "WEBP" if fmt == "webp" else "PNG"
    img.save(output_path, save_format)
    print(f"Transparent {save_format} saved to: {output_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Remove green screen from image")
    parser.add_argument("input", help="Input image path (with green screen)")
    parser.add_argument("output", help="Output PNG path (transparent)")
    parser.add_argument("--tolerance", type=int, default=60,
                        help="Green detection tolerance 0-255 (default: 60)")
    parser.add_argument("--edge-blur", type=int, default=1,
                        help="Edge anti-alias blur radius in pixels (default: 1)")
    parser.add_argument("--format", choices=["png", "webp"], default="png",
                        help="Output format: png or webp (default: png)")
    args = parser.parse_args()

    if not Path(args.input).exists():
        print(f"ERROR: Input file not found: {args.input}", file=sys.stderr)
        sys.exit(1)

    remove_green(args.input, args.output, args.tolerance, args.edge_blur, args.format)


if __name__ == "__main__":
    main()
