settled-reach/tooling/visual-diff

#!/usr/bin/env python3
"""Pixel-level visual diff for golden image comparison.

Compares two PNG images per-channel with configurable tolerance.
Reads default tolerance from tests/visual.json if available.

Usage:
  tooling/visual-diff EXPECTED ACTUAL [--tolerance N] [--diff-output PATH] [--config PATH]

Exit codes:
  0 = images match (all pixels within tolerance)
  1 = images differ
  2 = size mismatch or fatal error
"""

import argparse
import json
import struct
import sys
import zlib
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_CONFIG = ROOT / "tests" / "visual.json"
DEFAULT_TOLERANCE = 5

# ---------------------------------------------------------------------------
# PNG reading — prefer PIL, fallback to pure stdlib
# ---------------------------------------------------------------------------

_USE_PIL = False
try:
    from PIL import Image as _PILImage

    _USE_PIL = True
except ImportError:
    pass


def _read_png_pil(path: str) -> tuple[int, int, bytes]:
    """Read PNG via Pillow, return (width, height, RGBA bytes)."""
    img = _PILImage.open(path).convert("RGBA")
    return img.width, img.height, img.tobytes()


def _paeth(a: int, b: int, c: int) -> int:
    p = a + b - c
    pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
    if pa <= pb and pa <= pc:
        return a
    if pb <= pc:
        return b
    return c


def _read_png_stdlib(path: str) -> tuple[int, int, bytes]:
    """Read an RGBA (color type 6) PNG using only struct + zlib.

    Handles multiple IDAT chunks and all 5 PNG filter types.
    """
    with open(path, "rb") as f:
        sig = f.read(8)
        if sig != b"\x89PNG\r\n\x1a\n":
            print(f"ERROR: {path} is not a valid PNG", file=sys.stderr)
            sys.exit(2)

        width = height = 0
        bit_depth = color_type = 0
        idat_chunks: list[bytes] = []

        while True:
            header = f.read(8)
            if len(header) < 8:
                break
            length, chunk_type = struct.unpack(">I4s", header)
            data = f.read(length)
            _crc = f.read(4)

            if chunk_type == b"IHDR":
                width, height, bit_depth, color_type = struct.unpack(
                    ">IIBB", data[:10]
                )
                if color_type != 6:
                    print(
                        f"ERROR: {path} has color type {color_type}, expected 6 (RGBA)",
                        file=sys.stderr,
                    )
                    sys.exit(2)
                if bit_depth != 8:
                    print(
                        f"ERROR: {path} has bit depth {bit_depth}, expected 8",
                        file=sys.stderr,
                    )
                    sys.exit(2)
            elif chunk_type == b"IDAT":
                idat_chunks.append(data)
            elif chunk_type == b"IEND":
                break

    raw = zlib.decompress(b"".join(idat_chunks))

    bpp = 4  # RGBA = 4 bytes per pixel
    stride = width * bpp
    pixels = bytearray(height * stride)

    pos = 0
    for y in range(height):
        filter_type = raw[pos]
        pos += 1
        row_start = y * stride

        for x in range(stride):
            cur = raw[pos]
            pos += 1

            a = pixels[row_start + x - bpp] if x >= bpp else 0
            b = pixels[row_start - stride + x] if y > 0 else 0
            c = (
                pixels[row_start - stride + x - bpp]
                if y > 0 and x >= bpp
                else 0
            )

            if filter_type == 0:  # None
                val = cur
            elif filter_type == 1:  # Sub
                val = (cur + a) & 0xFF
            elif filter_type == 2:  # Up
                val = (cur + b) & 0xFF
            elif filter_type == 3:  # Average
                val = (cur + ((a + b) >> 1)) & 0xFF
            elif filter_type == 4:  # Paeth
                val = (cur + _paeth(a, b, c)) & 0xFF
            else:
                print(
                    f"ERROR: unknown PNG filter type {filter_type} at row {y}",
                    file=sys.stderr,
                )
                sys.exit(2)

            pixels[row_start + x] = val

    return width, height, bytes(pixels)


def read_png(path: str) -> tuple[int, int, bytes]:
    """Read PNG, return (width, height, RGBA bytes)."""
    if _USE_PIL:
        return _read_png_pil(path)
    return _read_png_stdlib(path)


# ---------------------------------------------------------------------------
# Diff PNG writing — prefer PIL, fallback to pure stdlib
# ---------------------------------------------------------------------------


def _write_png_pil(path: str, width: int, height: int, rgba: bytes) -> None:
    img = _PILImage.frombytes("RGBA", (width, height), rgba)
    img.save(path)


def _write_png_stdlib(
    path: str, width: int, height: int, rgba: bytes
) -> None:
    """Write a minimal RGBA PNG using zlib + struct (filter type 0/None)."""

    def _chunk(chunk_type: bytes, data: bytes) -> bytes:
        crc = zlib.crc32(chunk_type + data) & 0xFFFFFFFF
        return struct.pack(">I", len(data)) + chunk_type + data + struct.pack(">I", crc)

    # IHDR: width, height, bit_depth=8, color_type=6, compress=0, filter=0, interlace=0
    ihdr_data = struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0)

    # Build raw scanlines with filter byte 0 (None) per row
    stride = width * 4
    raw = bytearray()
    for y in range(height):
        raw.append(0)  # filter type None
        offset = y * stride
        raw.extend(rgba[offset : offset + stride])

    compressed = zlib.compress(bytes(raw))

    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(_chunk(b"IHDR", ihdr_data))
        f.write(_chunk(b"IDAT", compressed))
        f.write(_chunk(b"IEND", b""))


def write_png(path: str, width: int, height: int, rgba: bytes) -> None:
    if _USE_PIL:
        _write_png_pil(path, width, height, rgba)
    else:
        _write_png_stdlib(path, width, height, rgba)


# ---------------------------------------------------------------------------
# Comparison
# ---------------------------------------------------------------------------


def compare(
    expected: bytes,
    actual: bytes,
    width: int,
    height: int,
    tolerance: int,
) -> tuple[int, bytes | None]:
    """Compare two RGBA buffers. Returns (diff_count, diff_rgba_or_None)."""
    total = width * height
    diff_count = 0
    diff_buf = bytearray(total * 4)

    for i in range(total):
        off = i * 4
        er, eg, eb, ea = expected[off], expected[off + 1], expected[off + 2], expected[off + 3]
        ar, ag, ab, aa = actual[off], actual[off + 1], actual[off + 2], actual[off + 3]

        if (
            abs(er - ar) > tolerance
            or abs(eg - ag) > tolerance
            or abs(eb - ab) > tolerance
            or abs(ea - aa) > tolerance
        ):
            diff_count += 1
            diff_buf[off] = 0xFF
            diff_buf[off + 1] = 0x00
            diff_buf[off + 2] = 0x00
            diff_buf[off + 3] = 0xFF
        # else: remains (0, 0, 0, 0) — transparent

    return diff_count, bytes(diff_buf)


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


def load_tolerance(config_path: Path | None) -> int:
    """Read tolerance from config JSON, return DEFAULT_TOLERANCE on failure."""
    if config_path is None:
        config_path = DEFAULT_CONFIG
    if not config_path.exists():
        return DEFAULT_TOLERANCE
    try:
        with open(config_path) as f:
            data = json.load(f)
        return int(data.get("tolerance", DEFAULT_TOLERANCE))
    except (json.JSONDecodeError, ValueError, OSError):
        return DEFAULT_TOLERANCE


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Pixel-level visual diff for golden image comparison."
    )
    parser.add_argument("expected", help="Path to golden PNG")
    parser.add_argument("actual", help="Path to captured PNG")
    parser.add_argument(
        "--tolerance",
        type=int,
        default=None,
        help="Per-channel pixel tolerance (default: from config or 5)",
    )
    parser.add_argument(
        "--diff-output",
        default=None,
        help="Path to write diff PNG highlighting changed pixels",
    )
    parser.add_argument(
        "--config",
        default=None,
        help="Path to tests/visual.json (default: auto-detect)",
    )
    args = parser.parse_args()

    # Resolve tolerance: CLI > config > fallback
    config_path = Path(args.config) if args.config else None
    tolerance = args.tolerance if args.tolerance is not None else load_tolerance(config_path)

    # Read images
    try:
        ew, eh, epx = read_png(args.expected)
    except FileNotFoundError:
        print(f"ERROR: expected image not found: {args.expected}", file=sys.stderr)
        return 2
    except Exception as exc:
        print(f"ERROR: failed to read expected image: {exc}", file=sys.stderr)
        return 2

    try:
        aw, ah, apx = read_png(args.actual)
    except FileNotFoundError:
        print(f"ERROR: actual image not found: {args.actual}", file=sys.stderr)
        return 2
    except Exception as exc:
        print(f"ERROR: failed to read actual image: {exc}", file=sys.stderr)
        return 2

    # Size check
    if ew != aw or eh != ah:
        print(
            f"ERROR: size mismatch — expected {ew}x{eh}, actual {aw}x{ah}",
            file=sys.stderr,
        )
        return 2

    # Compare
    diff_count, diff_buf = compare(epx, apx, ew, eh, tolerance)
    total = ew * eh

    if diff_count == 0:
        print(f"PASS: images match ({ew}x{eh})")
        return 0

    pct = diff_count / total * 100
    print(f"FAIL: {diff_count} of {total} pixels differ ({pct:.1f}%)")

    if args.diff_output and diff_buf:
        Path(args.diff_output).parent.mkdir(parents=True, exist_ok=True)
        write_png(args.diff_output, ew, eh, diff_buf)

    return 1


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