#!/usr/bin/env python3
"""
compare.py  <baseline.png> <candidate.png>

Pixel-level comparison of two PNG screenshots.

Exit codes:
  0  images are identical (or within tolerance)
  1  images differ beyond tolerance
  2  usage / file error

Outputs a one-line summary and, on failure, saves a diff image alongside
the candidate as <candidate>_diff.png highlighting changed pixels in red.
"""

import struct
import sys
import zlib

TOLERANCE  = 10    # per-channel delta considered noise
THRESHOLD  = 0.01  # max fraction of differing pixels before failure (1 %)


# ── Minimal PNG reader ────────────────────────────────────────────────────────

def _read_chunk(data, pos):
    length = struct.unpack_from('>I', data, pos)[0]
    tag    = data[pos+4:pos+8]
    body   = data[pos+8:pos+8+length]
    pos   += 12 + length
    return tag, body, pos

def load_png(path):
    with open(path, 'rb') as f:
        raw = f.read()
    assert raw[:8] == b'\x89PNG\r\n\x1a\n', f"{path}: not a PNG"

    ihdr = width = height = None
    idat_chunks = []
    pos = 8
    while pos < len(raw):
        tag, body, pos = _read_chunk(raw, pos)
        if tag == b'IHDR':
            width, height = struct.unpack('>II', body[:8])
            bit_depth, colour_type = body[8], body[9]
            assert bit_depth == 8 and colour_type == 2, \
                f"{path}: only 8-bit RGB PNGs supported"
        elif tag == b'IDAT':
            idat_chunks.append(body)
        elif tag == b'IEND':
            break

    assert width and height, f"{path}: no IHDR"
    raw_pixels = zlib.decompress(b''.join(idat_chunks))

    # Reconstruct pixels via PNG filter (only filter 0 = None is used here)
    stride = width * 3
    pixels = []
    idx = 0
    prev_row = bytes(stride)
    for _ in range(height):
        filt = raw_pixels[idx]; idx += 1
        row  = bytearray(raw_pixels[idx:idx+stride]); idx += stride
        if filt == 0:
            pass
        elif filt == 1:  # Sub
            for i in range(3, stride):
                row[i] = (row[i] + row[i-3]) & 0xFF
        elif filt == 2:  # Up
            for i in range(stride):
                row[i] = (row[i] + prev_row[i]) & 0xFF
        elif filt == 3:  # Average
            for i in range(stride):
                a = row[i-3] if i >= 3 else 0
                row[i] = (row[i] + (a + prev_row[i]) // 2) & 0xFF
        elif filt == 4:  # Paeth
            for i in range(stride):
                a = row[i-3]   if i >= 3 else 0
                b = prev_row[i]
                c = prev_row[i-3] if i >= 3 else 0
                p  = a + b - c
                pa = abs(p - a); pb = abs(p - b); pc = abs(p - c)
                pr = a if pa <= pb and pa <= pc else (b if pb <= pc else c)
                row[i] = (row[i] + pr) & 0xFF
        prev_row = bytes(row)
        pixels.append([tuple(row[i:i+3]) for i in range(0, stride, 3)])

    return width, height, pixels


def save_png(path, pixels):
    h = len(pixels)
    w = len(pixels[0])
    rows = bytearray()
    for row in pixels:
        rows.append(0)  # filter = None
        for p in row:
            rows += bytearray(p)

    def chunk(tag, data):
        crc = zlib.crc32(tag + data) & 0xffffffff
        return struct.pack('>I', len(data)) + tag + data + struct.pack('>I', crc)

    out  = b'\x89PNG\r\n\x1a\n'
    out += chunk(b'IHDR', struct.pack('>IIBBBBB', w, h, 8, 2, 0, 0, 0))
    out += chunk(b'IDAT', zlib.compress(bytes(rows), 9))
    out += chunk(b'IEND', b'')
    with open(path, 'wb') as f:
        f.write(out)


# ── Comparison logic ──────────────────────────────────────────────────────────

def compare(baseline_path, candidate_path):
    try:
        bw, bh, base_px = load_png(baseline_path)
        cw, ch, cand_px = load_png(candidate_path)
    except Exception as e:
        print(f"ERROR: {e}")
        return 2

    if bw != cw or bh != ch:
        print(f"FAIL  size mismatch: baseline={bw}×{bh}  candidate={cw}×{ch}")
        return 1

    total    = bw * bh
    diff_px  = 0
    diff_img = [[(0, 0, 0)] * bw for _ in range(bh)]

    for y in range(bh):
        for x in range(bw):
            br, bg, bb = base_px[y][x]
            cr, cg, cb = cand_px[y][x]
            if abs(br-cr) > TOLERANCE or abs(bg-cg) > TOLERANCE or abs(bb-cb) > TOLERANCE:
                diff_px += 1
                diff_img[y][x] = (220, 0, 0)   # red highlight on diff image
            else:
                # Dim the matching pixel so diffs stand out
                diff_img[y][x] = (cr//3, cg//3, cb//3)

    frac = diff_px / total

    if frac > THRESHOLD:
        diff_path = candidate_path.replace('.png', '_diff.png')
        save_png(diff_path, diff_img)
        print(f"FAIL  {diff_px}/{total} pixels differ ({frac:.2%}) "
              f"— diff saved to {diff_path}")
        return 1
    else:
        print(f"PASS  {diff_px}/{total} pixels differ ({frac:.2%}) "
              f"[tolerance ≤{THRESHOLD:.0%}]")
        return 0


if __name__ == '__main__':
    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} <baseline.png> <candidate.png>")
        sys.exit(2)
    sys.exit(compare(sys.argv[1], sys.argv[2]))