#!/usr/bin/env python3
"""Slice a fixed grid, remove green chroma, and resize with premultiplied alpha."""

from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("image", type=Path)
    parser.add_argument("--rows", type=int, required=True)
    parser.add_argument("--cols", type=int, required=True)
    parser.add_argument("--names", required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--inset", type=float, default=0.94)
    return parser.parse_args()


def remove_green(image: Image.Image) -> Image.Image:
    rgba = image.convert("RGBA")
    cleaned: list[tuple[int, int, int, int]] = []
    for red, green, blue, alpha in rgba.getdata():
        key_distance = max(red, 255 - green, blue)
        green_dominant = green - red > 45 and green - blue > 45
        if key_distance <= 90 or green_dominant:
            cleaned.append((0, 0, 0, 0))
        else:
            cleaned.append((red, green, blue, alpha))
    rgba.putdata(cleaned)
    return rgba


def resize_alpha(image: Image.Image, pixels: int) -> Image.Image:
    # Premultiplied resize prevents hidden chroma RGB from leaking through the
    # Lanczos filter into otherwise clean visible edge pixels.
    resized = image.convert("RGBa").resize(
        (pixels, pixels), Image.Resampling.LANCZOS
    ).convert("RGBA")
    cleaned: list[tuple[int, int, int, int]] = []
    for red, green, blue, alpha in resized.getdata():
        if green - red > 45 and green - blue > 45:
            cleaned.append((0, 0, 0, 0))
        else:
            cleaned.append((red, green, blue, alpha))
    resized.putdata(cleaned)
    return resized


def main() -> None:
    args = parse_args()
    if args.rows < 1 or args.cols < 1 or not 0 < args.inset <= 1:
        raise SystemExit("invalid grid or inset")
    names = [name.strip() for name in args.names.split(",") if name.strip()]
    if len(names) != args.rows * args.cols:
        raise SystemExit("name count must equal rows * cols")

    source = Image.open(args.image).convert("RGBA")
    width, height = source.size
    args.output_dir.mkdir(parents=True, exist_ok=True)
    for index, name in enumerate(names):
        row, col = divmod(index, args.cols)
        left = round(col * width / args.cols)
        right = round((col + 1) * width / args.cols)
        top = round(row * height / args.rows)
        bottom = round((row + 1) * height / args.rows)
        cell_width, cell_height = right - left, bottom - top
        crop_width = round(cell_width * args.inset)
        crop_height = round(cell_height * args.inset)
        center_x, center_y = (left + right) / 2, (top + bottom) / 2
        box = (
            round(center_x - crop_width / 2),
            round(center_y - crop_height / 2),
            round(center_x + crop_width / 2),
            round(center_y + crop_height / 2),
        )
        asset = remove_green(source.crop(box))
        resize_alpha(asset, 128).save(args.output_dir / f"{name}@1x.png")
        resize_alpha(asset, 256).save(args.output_dir / f"{name}@2x.png")

    print(f"GREEN_GRID_PROCESS_OK assets={len(names) * 2} output={args.output_dir}")


if __name__ == "__main__":
    main()