#!/usr/bin/env python3
"""Compares two folders of screenshots from tools/take_screenshots.gd pixel by pixel.

Usage: tools/compare_screenshots.py <before dir> <after dir> [<diff dir>]

Prints how many pixels differ in each screenshot. With a diff dir, writes an image
for each screenshot that differs, with the changed pixels in red over a faded copy of
the after picture. Exits with 1 if any screenshot differs or is missing.
"""

import sys
from pathlib import Path

from PIL import Image, ImageChops


def main() -> int:
    if len(sys.argv) < 3:
        print(__doc__)
        return 2
    before_dir, after_dir = Path(sys.argv[1]), Path(sys.argv[2])
    diff_dir = Path(sys.argv[3]) if len(sys.argv) > 3 else None
    if diff_dir:
        diff_dir.mkdir(parents=True, exist_ok=True)

    names = sorted({p.name for p in before_dir.glob("*.png")}
                   | {p.name for p in after_dir.glob("*.png")})
    failed = False
    for name in names:
        before_path, after_path = before_dir / name, after_dir / name
        if not before_path.exists() or not after_path.exists():
            print(f"{name}: missing in {'before' if not before_path.exists() else 'after'}")
            failed = True
            continue
        before = Image.open(before_path).convert("RGB")
        after = Image.open(after_path).convert("RGB")
        if before.size != after.size:
            print(f"{name}: size {before.size} -> {after.size}")
            failed = True
            continue
        mask = ImageChops.difference(before, after).convert("L").point(
            lambda value: 255 if value else 0)
        changed = mask.histogram()[255]
        print(f"{name}: {changed} pixels differ")
        if changed:
            failed = True
            if diff_dir:
                faded = Image.blend(after, Image.new("RGB", after.size, "white"), 0.7)
                faded.paste(Image.new("RGB", after.size, "red"), mask=mask)
                faded.save(diff_dir / name)
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())