#!/usr/bin/env python3
"""Pixel-exact comparison of a rendered snapshot against a checked-in reference.

CI gate: the smoke job renders deterministic frames (see `just gen-check`)
and fails if they differ from docs/images/reference-*.png. The references are
regenerated by scripts/gen-media.py (`just gen`) — a PR that
intentionally changes the office's look must commit the regenerated references
in the same change.

Exit codes: 0 = identical, 1 = images differ (diff overlay written),
2 = could not compare (usage / missing / unreadable / size mismatch).

NOTE: compare in RGB, never RGBA — ImageChops.difference of two opaque RGBA
images has an all-zero alpha band, and Pillow >= 10's getbbox() defaults to
alpha_only=True, which would make ANY two same-size images compare equal.
"""

import sys
from pathlib import Path

from PIL import Image, ImageChops, UnidentifiedImageError


def compare_images(ref_path, cand_path, diff_path, threshold_percent=0.0):
    """Returns an exit code (see module docstring)."""
    ref_path = Path(ref_path)
    cand_path = Path(cand_path)
    diff_path = Path(diff_path)

    try:
        ref = Image.open(ref_path).convert("RGB")
    except (OSError, UnidentifiedImageError) as e:
        print(f"ERROR: cannot read reference image {ref_path}: {e}")
        return 2
    try:
        cand = Image.open(cand_path).convert("RGB")
    except (OSError, UnidentifiedImageError) as e:
        print(f"ERROR: cannot read candidate image {cand_path}: {e}")
        return 2

    if ref.size != cand.size:
        print(f"ERROR: Image sizes differ. Reference: {ref.size}, Candidate: {cand.size}")
        return 2

    diff = ImageChops.difference(ref, cand)
    if diff.getbbox() is None:
        print(f"SUCCESS: {cand_path.name} is pixel-identical to {ref_path.name}.")
        return 0

    # Generate a visual difference overlay highlighting mismatches in red.
    diff_data = diff.load()
    cand_data = cand.load()

    w, h = ref.size
    mismatched = 0

    out = Image.new("RGB", (w, h))
    out_data = out.load()

    for y in range(h):
        for x in range(w):
            if diff_data[x, y] != (0, 0, 0):
                mismatched += 1
                out_data[x, y] = (255, 0, 0)  # highlight mismatched pixels in red
            else:
                # Dim the matching background so the red overlay pops.
                c_pix = cand_data[x, y]
                out_data[x, y] = (c_pix[0] // 2, c_pix[1] // 2, c_pix[2] // 2)

    total_pixels = w * h
    percent = (mismatched / total_pixels) * 100.0
    print(f"MISMATCH: {mismatched} / {total_pixels} pixels differ ({percent:.4f}%)")

    diff_path.parent.mkdir(parents=True, exist_ok=True)
    out.save(diff_path)
    print(f"Difference highlight image saved to: {diff_path}")

    # The render is byte-deterministic, so the default tolerance is zero; a
    # nonzero threshold must be an explicit, visible choice at the call site.
    if percent > threshold_percent:
        print(f"ERROR: Mismatch {percent:.4f}% exceeds the allowed threshold of {threshold_percent}%.")
        return 1

    print("SUCCESS: Mismatch is within the acceptable threshold.")
    return 0


if __name__ == "__main__":
    if len(sys.argv) not in (4, 5):
        print("Usage: compare-screenshots.py <reference> <candidate> <diff_out> [threshold_percent]")
        sys.exit(2)

    threshold = float(sys.argv[4]) if len(sys.argv) == 5 else 0.0
    sys.exit(compare_images(sys.argv[1], sys.argv[2], sys.argv[3], threshold))
