#!/usr/bin/env python3

from pathlib import Path
from PIL import Image, ImageChops
import argparse
import re

def load_rgb(path: Path) -> Image.Image:
    return Image.open(path).convert("RGB")

def mse(img: Image.Image) -> float:
    hist = img.histogram()
    sq = sum(v * ((i % 256) ** 2) for i, v in enumerate(hist))
    return sq / (img.width * img.height * 3)

def crop_pair(expected, actual, shift):
    # shift > 0 means actual content is lower than expected.
    w = min(expected.width, actual.width)

    if shift >= 0:
        e_box = (0, 0, w, min(expected.height, actual.height - shift))
        a_box = (0, shift, w, shift + (e_box[3] - e_box[1]))
    else:
        s = -shift
        a_box = (0, 0, w, min(actual.height, expected.height - s))
        e_box = (0, s, w, s + (a_box[3] - a_box[1]))

    return expected.crop(e_box), actual.crop(a_box)

def best_vertical_shift(expected, actual, max_shift=100):
    best = None

    for shift in range(18, 19):
        e, a = crop_pair(expected, actual, shift)
        if e.height <= 0:
            continue

        score = mse(ImageChops.difference(e, a))
        if best is None or score < best[1]:
            best = (shift, score)

    return best

def changed_until_y(expected, actual, threshold=5, min_changed_pixels=10):
    diff = ImageChops.difference(expected, actual)
    last_changed = None

    for y in range(diff.height):
        row = diff.crop((0, y, diff.width, y + 1)).convert("L")
        changed = sum(1 for px in row.getdata() if px > threshold)

        if changed >= min_changed_pixels:
            last_changed = y

    return {
        "last_changed_y": last_changed,
        "clean_after_y": 0 if last_changed is None else last_changed + 1,
    }

def find_pairs(root: Path):
    expected_files = sorted(root.rglob("*-expected.png"))

    for expected in expected_files:
        actual = expected.with_name(
            expected.name.replace("-expected.png", "-actual.png")
        )
        if actual.exists():
            yield expected, actual

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("root", type=Path, nargs="?", default=Path("test-results"))
    ap.add_argument("--max-shift", type=int, default=100)
    ap.add_argument("--make-diffs", action="store_true")
    args = ap.parse_args()

    failed = False

    for expected_path, actual_path in find_pairs(args.root):
        expected = load_rgb(expected_path)
        actual = load_rgb(actual_path)

        raw_score = mse(ImageChops.difference(
            expected.crop((0, 0, min(expected.width, actual.width), min(expected.height, actual.height))),
            actual.crop((0, 0, min(expected.width, actual.width), min(expected.height, actual.height))),
        ))

        shift, shifted_score = best_vertical_shift(
            expected, actual, max_shift=args.max_shift
        )

        status = "OK" if shifted_score < raw_score * 0.05 else "CHECK"
        if status == "CHECK":
            failed = True

        rel = expected_path.relative_to(args.root)
        print(f"{status} {rel}")
        print(f"  expected: {expected.size}, actual: {actual.size}")
        print(f"  best shift: {shift:+d}px")
        print(f"  raw mse:     {raw_score:.3f}")
        print(f"  shifted mse: {shifted_score:.3f}")

        if args.make_diffs:
            e, a = crop_pair(expected, actual, shift)
            change_info = changed_until_y(e, a, threshold=2)
            print(f"  last changed y after shift: {change_info['last_changed_y']}")
            print(f"  clean after y after shift:  {change_info['clean_after_y']}")
            diff = ImageChops.difference(e, a)
            out = expected_path.with_name(
                expected_path.name.replace("-expected.png", f"-shift-{shift:+d}-diff.png")
            )
            diff.save(out)
            print(f"  wrote: {out}")

    raise SystemExit(1 if failed else 0)

if __name__ == "__main__":
    main()