settled-reach/tooling/domains/visual/diff.py

#!/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:
  reach 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 json
import struct
import zlib
from pathlib import Path

from tooling.core import config, console
from tooling.core.errors import ReachError

ROOT = config.repo_root()
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":
            console.event(f"ERROR: {path} is not a valid PNG", level="error")
            raise ReachError(
                "visual-diff: unreadable PNG",
                fix="the file is not an 8-bit RGBA PNG — re-capture it",
                exit_code=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:
                    console.event(
                        f"ERROR: {path} has color type {color_type}, expected 6 (RGBA)",
                        level="error",
                    )
                    raise ReachError(
                "visual-diff: unreadable PNG",
                fix="the file is not an 8-bit RGBA PNG — re-capture it",
                exit_code=2,
            )
                if bit_depth != 8:
                    console.event(
                        f"ERROR: {path} has bit depth {bit_depth}, expected 8",
                        level="error",
                    )
                    raise ReachError(
                "visual-diff: unreadable PNG",
                fix="the file is not an 8-bit RGBA PNG — re-capture it",
                exit_code=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:
                console.event(
                    f"ERROR: unknown PNG filter type {filter_type} at row {y}",
                    level="error",
                )
                raise ReachError(
                "visual-diff: unreadable PNG",
                fix="the file is not an 8-bit RGBA PNG — re-capture it",
                exit_code=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_config(config_path: Path | None) -> dict:
    """Read visual test config, return dict with tolerance and max_diff_pct."""
    if config_path is None:
        config_path = DEFAULT_CONFIG
    defaults = {"tolerance": DEFAULT_TOLERANCE, "max_diff_pct": 0.0}
    if not config_path.exists():
        return defaults
    try:
        with open(config_path) as f:
            data = json.load(f)
        return {
            "tolerance": int(data.get("tolerance", DEFAULT_TOLERANCE)),
            "max_diff_pct": float(data.get("max_diff_pct", 0.0)),
        }
    except (json.JSONDecodeError, ValueError, OSError):
        return defaults


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


def run(
    expected: str,
    actual: str,
    tolerance: int | None = None,
    max_diff_pct: float | None = None,
    diff_output: str | None = None,
    config_path_arg: str | None = None,
) -> int:
    """Pixel-level visual diff. Returns 0 pass, 1 fail, 2 unusable input.

    The argparse parser that used to live here is gone: the router declares the
    options, and a parser inside a service would be a second transport layer
    (D-263)."""

    # Resolve settings: CLI > config > fallback
    config_path = Path(config_path_arg) if config_path_arg else None
    cfg = load_config(config_path)
    tolerance = tolerance if tolerance is not None else cfg["tolerance"]
    max_diff_pct = max_diff_pct if max_diff_pct is not None else cfg["max_diff_pct"]

    # Read images
    try:
        ew, eh, epx = read_png(expected)
    except FileNotFoundError:
        console.event(f"ERROR: expected image not found: {expected}", level="error")
        return 2
    except Exception as exc:
        console.event(f"ERROR: failed to read expected image: {exc}", level="error")
        return 2

    try:
        aw, ah, apx = read_png(actual)
    except FileNotFoundError:
        console.event(f"ERROR: actual image not found: {actual}", level="error")
        return 2
    except Exception as exc:
        console.event(f"ERROR: failed to read actual image: {exc}", level="error")
        return 2

    # Size check
    if ew != aw or eh != ah:
        console.event(
            f"ERROR: size mismatch — expected {ew}x{eh}, actual {aw}x{ah}",
            level="error",
        )
        return 2

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

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

    pct = diff_count / total * 100

    if pct <= max_diff_pct:
        console.event(f"PASS: {diff_count} of {total} pixels differ ({pct:.1f}%, within {max_diff_pct}% threshold)")
        return 0

    console.event(f"FAIL: {diff_count} of {total} pixels differ ({pct:.1f}%)")

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

    return 1