from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image, ImageChops


ICON_SIZES = (16, 20, 24, 32, 40, 48, 64, 128, 256)


def remove_green_background(image: Image.Image) -> Image.Image:
	rgba = image.convert("RGBA")
	pixels = rgba.load()
	width, height = rgba.size

	for y in range(height):
		for x in range(width):
			r, g, b, a = pixels[x, y]
			green_score = g - max(r, b)
			is_key = g > 90 and green_score > 35
			is_fringe = g > 32 and g > r * 1.35 and g > b * 1.20
			if is_key or is_fringe:
				edge = max(0, min(255, (green_score - 18) * 5))
				alpha = max(0, 255 - edge)
				if alpha < 56:
					pixels[x, y] = (0, 0, 0, 0)
				else:
					despilled_g = min(g, max(r, b))
					pixels[x, y] = (r, despilled_g, b, min(a, alpha))

	return rgba


def crop_to_alpha(image: Image.Image, padding_ratio: float) -> Image.Image:
    alpha = image.getchannel("A")
    bbox = alpha.getbbox()
    if not bbox:
        raise ValueError("source image has no visible pixels after background removal")

    cropped = image.crop(bbox)
    side = max(cropped.size)
    padding = round(side * padding_ratio)
    canvas_side = side + padding * 2
    canvas = Image.new("RGBA", (canvas_side, canvas_side), (0, 0, 0, 0))
    canvas.alpha_composite(cropped, ((canvas_side - cropped.width) // 2, (canvas_side - cropped.height) // 2))
    return canvas


def trim_transparent_padding(image: Image.Image) -> Image.Image:
    empty = Image.new("RGBA", image.size, (0, 0, 0, 0))
    bbox = ImageChops.difference(image, empty).getbbox()
    if bbox is None:
        return image
    return image.crop(bbox)


def build_icons(source: Path, out_dir: Path, padding_ratio: float) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)

    source_image = Image.open(source)
    transparent = remove_green_background(source_image)
    master = crop_to_alpha(transparent, padding_ratio)
    master = trim_transparent_padding(master)

    master_path = out_dir / "tray-icon.png"
    master.save(master_path)

    resized_images: list[Image.Image] = []
    for size in ICON_SIZES:
        resized = master.resize((size, size), Image.Resampling.LANCZOS)
        resized.save(out_dir / f"tray-icon-{size}.png")
        resized_images.append(resized)

    ico_path = out_dir / "tray.ico"
    resized_images[-1].save(
        ico_path,
        format="ICO",
        sizes=[(size, size) for size in ICON_SIZES],
        append_images=resized_images[:-1],
    )

    print(f"wrote {master_path}")
    print(f"wrote {ico_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Build u-desktop tray icon assets from a generated source image.")
    parser.add_argument("--source", type=Path, required=True, help="Generated source PNG on a green chroma-key background.")
    parser.add_argument("--out-dir", type=Path, default=Path("assets/icons"), help="Directory for PNG sizes and tray.ico.")
    parser.add_argument("--padding", type=float, default=0.04, help="Transparent padding ratio around the cropped icon.")
    args = parser.parse_args()

    build_icons(args.source, args.out_dir, args.padding)


if __name__ == "__main__":
    main()