from pathlib import Path

from PIL import Image, ImageChops, ImageFilter, ImageStat
from PIL.PngImagePlugin import PngInfo


ROOT = Path(__file__).resolve().parents[1]
UNIT_DIR = ROOT / "src" / "assets" / "images" / "units"
PROCESSING_VERSION = "solid-alpha-silhouette-v1"


def solid_alpha(value: int) -> int:
    return 0 if value < 16 else 255


def unit_fill_color(image: Image.Image, mask: Image.Image) -> tuple[int, int, int, int]:
    stat = ImageStat.Stat(image.convert("RGB"), mask=mask)
    if not stat.count[0]:
        return (54, 45, 34, 255)

    red, green, blue = stat.mean
    return (max(30, round(red * 0.45)), max(26, round(green * 0.45)), max(22, round(blue * 0.45)), 255)


def solidify_image(image: Image.Image) -> Image.Image:
    image = image.convert("RGBA")
    red, green, blue, alpha = image.split()
    fixed_alpha = alpha.point(solid_alpha)
    solid = Image.merge("RGBA", (red, green, blue, fixed_alpha))
    transparent = Image.new("RGBA", image.size, (0, 0, 0, 0))
    solid = Image.composite(solid, transparent, fixed_alpha)

    silhouette = fixed_alpha.filter(ImageFilter.MaxFilter(15)).filter(ImageFilter.MinFilter(7))
    silhouette = silhouette.point(lambda value: 255 if value >= 255 else 0)
    fill_mask = ImageChops.subtract(silhouette, fixed_alpha)
    fill = Image.new("RGBA", image.size, unit_fill_color(image, fixed_alpha))
    body = Image.composite(fill, transparent, fill_mask)
    return Image.alpha_composite(body, solid)


def solidify_unit_alpha(path: Path) -> bool:
    with Image.open(path) as source:
        if source.info.get("heros_unit_alpha") == PROCESSING_VERSION:
            return False
        image = source.convert("RGBA")
    fixed = solidify_image(image)

    if ImageChops.difference(image, fixed).getbbox() is None:
        return False

    png_info = PngInfo()
    png_info.add_text("heros_unit_alpha", PROCESSING_VERSION)
    fixed.save(path, optimize=True, pnginfo=png_info)
    return True


def main() -> None:
    changed = 0
    for path in sorted(UNIT_DIR.glob("unit-*.png")):
        if solidify_unit_alpha(path):
            changed += 1

    print(f"Solidified alpha in {changed} unit sprite sheets.")


if __name__ == "__main__":
    main()