#!/usr/bin/env python3
# Author:  cirakg & claude ai
# Summary: Cut the watch screen out of a simulator window screenshot and compare it with the approved golden image.
"""Cut the watch screen out of a Connect IQ simulator screenshot.

Usage:
    sim_screenshot.py --raw RAW.png --device-dir DIR --out SCREEN.png
                      [--golden GOLDEN.png] [--approve]

Steps:
 1. Find where the device drawing (DIR/device.png) sits in the window
    screenshot. The window adds a menu bar and borders, so the offset is
    searched for, comparing only the bezel: the display area is masked out,
    because the app draws there.
 2. Crop the display rectangle from DIR/simulator.json (display.location) and
    black out the corners of round displays, which show bezel, not screen.
 3. Fail if the screen is blank (the app drew nothing).
 4. With --golden, compare with the approved image: the check fails when more
    than 0.5 % of the pixels differ by more than 24 levels in any channel. A
    diff image (differences in red) is written next to --out.
 5. With --approve, save the screen as the new golden image instead.

Exit: 0 ok, 1 differs from the golden image, 2 error, 3 no golden image yet.
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from pathlib import Path

from PIL import Image, ImageChops, ImageDraw, ImageStat
from PIL.PngImagePlugin import PngInfo

EXIT_OK, EXIT_DIFFERENT, EXIT_ERROR, EXIT_NO_GOLDEN = 0, 1, 2, 3
AUTHOR = "cirakg & claude ai"
MAX_OFFSET_X, MAX_OFFSET_Y = 16, 64  # search range for the device drawing
MIN_OVERLAP = 0.9  # the window may clip the drawing's edge, but not more than this
MAX_BEZEL_DIFF = 6.0  # mean absolute difference per channel at the best offset
PIXEL_TOLERANCE = 24  # per-channel difference that still counts as equal
MAX_DIFF_FRACTION = 0.005  # share of pixels allowed to differ
MIN_SCREEN_STDDEV = 2.0  # below this the screen counts as blank


@dataclass(frozen=True)
class Display:
    """The watch screen inside the device drawing."""

    x: int
    y: int
    width: int
    height: int
    shape: str  # round, rectangle, semi-octagon, ...
    image: str  # file name of the device drawing, e.g. device.png or fr255s.png


def open_image(path: Path) -> Image.Image:
    """Load an image completely and close its file straight away (no open handles left behind)."""
    with Image.open(path) as img:
        img.load()
        return img


def load_display(device_dir: Path) -> Display:
    """Display rectangle, shape and drawing file from the device's simulator.json."""
    sim = json.loads((device_dir / "simulator.json").read_text(encoding="utf-8"))
    loc = sim["display"]["location"]
    return Display(
        int(loc["x"]),
        int(loc["y"]),
        int(loc["width"]),
        int(loc["height"]),
        str(sim["display"]["shape"]),
        str(sim.get("image", "device.png")),
    )


def bezel_only(device: Image.Image, display: Display) -> Image.Image:
    """The device drawing with the display area made transparent."""
    masked = device.convert("RGBA")
    ImageDraw.Draw(masked).rectangle(
        (display.x, display.y, display.x + display.width - 1, display.y + display.height - 1), fill=(0, 0, 0, 0)
    )
    return masked


def find_offset(raw: Image.Image, bezel: Image.Image) -> tuple[int, int, float]:
    """Best (dx, dy) placement of the bezel in the screenshot, and its mean difference.

    When the window is exactly as wide as the drawing, the 1 px window border
    pushes the drawing's last column out of the window, so only the part that
    overlaps the screenshot is compared.
    """
    raw_rgba = raw.convert("RGBA")
    best = (0, 0, float("inf"))
    for dy in range(MAX_OFFSET_Y + 1):
        for dx in range(MAX_OFFSET_X + 1):
            width = min(bezel.width, raw.width - dx)
            height = min(bezel.height, raw.height - dy)
            if width < bezel.width * MIN_OVERLAP or height < bezel.height * MIN_OVERLAP:
                continue
            region = raw_rgba.crop((dx, dy, dx + width, dy + height))
            expected = region.copy()
            expected.alpha_composite(bezel.crop((0, 0, width, height)))
            diff = ImageChops.difference(expected.convert("RGB"), region.convert("RGB"))
            score = sum(ImageStat.Stat(diff).mean) / 3
            if score < best[2]:
                best = (dx, dy, score)
    return best


def crop_screen(raw: Image.Image, offset: tuple[int, int], display: Display) -> Image.Image:
    """The display area, with the corners of round screens blacked out."""
    dx, dy = offset
    left, top = dx + display.x, dy + display.y
    screen = raw.convert("RGB").crop((left, top, left + display.width, top + display.height))
    if display.shape == "round":
        mask = Image.new("L", screen.size, 0)
        ImageDraw.Draw(mask).ellipse((0, 0, display.width - 1, display.height - 1), fill=255)
        black = Image.new("RGB", screen.size, (0, 0, 0))
        screen = Image.composite(screen, black, mask)
    return screen


def is_blank(screen: Image.Image) -> bool:
    return max(ImageStat.Stat(screen).stddev) < MIN_SCREEN_STDDEV


def compare(screen: Image.Image, golden: Image.Image) -> tuple[float, Image.Image]:
    """Share of differing pixels and a diff image (differences in red)."""
    if screen.size != golden.size:
        return 1.0, Image.new("RGB", screen.size, (255, 0, 0))
    diff = ImageChops.difference(screen.convert("RGB"), golden.convert("RGB"))
    # A pixel differs when any channel differs by more than the tolerance.
    channels = [band.point(lambda v: 255 if v > PIXEL_TOLERANCE else 0) for band in diff.split()]
    differing = ImageChops.lighter(ImageChops.lighter(channels[0], channels[1]), channels[2])
    count = sum(differing.histogram()[1:])  # every non-zero pixel
    overlay = Image.composite(Image.new("RGB", screen.size, (255, 0, 0)), screen.convert("RGB"), differing)
    return count / (screen.width * screen.height), overlay


def save_png(img: Image.Image, path: Path, description: str) -> None:
    info = PngInfo()
    info.add_text("Author", AUTHOR)
    info.add_text("Description", description)
    path.parent.mkdir(parents=True, exist_ok=True)
    img.save(path, format="PNG", optimize=True, pnginfo=info)


class ScreenshotError(Exception):
    """A problem that makes the screenshot unusable; the message is shown to the user."""


def extract_screen(raw_path: Path, device_dir: Path, out: Path) -> Image.Image:
    """Find, crop, save and sanity-check the screen; raises ScreenshotError."""
    name = device_dir.name
    try:
        raw = open_image(raw_path)
        display = load_display(device_dir)
        device = open_image(device_dir / display.image)
    except (OSError, KeyError, ValueError) as err:
        raise ScreenshotError(f"cannot read the screenshot or the device files: {err}") from err
    dx, dy, score = find_offset(raw, bezel_only(device, display))
    if score > MAX_BEZEL_DIFF:
        raise ScreenshotError(f"device drawing not found in the screenshot (best difference {score:.1f})")
    print(f"[{name}] device drawing found at window offset ({dx}, {dy}), bezel difference {score:.2f}")
    screen = crop_screen(raw, (dx, dy), display)
    save_png(screen, out, f"Simulator screen of {name}")
    if is_blank(screen):
        raise ScreenshotError("the screen is blank: the app drew nothing")
    return screen


def check_golden(screen: Image.Image, name: str, out: Path, golden: Path) -> int:
    """Compare with the approved image; writes a diff image when they differ."""
    if not golden.exists():
        print(f"[{name}] no golden image yet: review {out}, then run 'make screenshots-approve'")
        return EXIT_NO_GOLDEN
    fraction, overlay = compare(screen, open_image(golden))
    if fraction > MAX_DIFF_FRACTION:
        diff_path = out.with_name(out.stem + "-diff.png")
        save_png(overlay, diff_path, f"Differences from the golden image on {name}")
        print(f"[{name}] differs from the golden image: {fraction:.2%} of pixels (see {diff_path})")
        return EXIT_DIFFERENT
    print(f"[{name}] matches the golden image ({fraction:.2%} of pixels differ)")
    return EXIT_OK


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Crop and check a simulator screenshot.")
    parser.add_argument("--raw", type=Path, required=True, help="screenshot of the simulator window")
    parser.add_argument("--device-dir", type=Path, required=True, help="the device's folder in ~/.Garmin")
    parser.add_argument("--out", type=Path, required=True, help="where to write the cropped screen")
    parser.add_argument("--golden", type=Path, help="approved image to compare with (or to write with --approve)")
    parser.add_argument("--approve", action="store_true", help="save the screen as the new golden image")
    args = parser.parse_args(argv)
    name = args.device_dir.name

    try:
        screen = extract_screen(args.raw, args.device_dir, args.out)
    except ScreenshotError as err:
        print(f"[{name}] {err}")
        return EXIT_ERROR
    if args.golden is None:
        print(f"[{name}] screen saved to {args.out}")
        return EXIT_OK
    if args.approve:
        save_png(screen, args.golden, f"Approved simulator screen of {name} (golden image for make screenshots)")
        print(f"[{name}] approved as {args.golden}")
        return EXIT_OK
    return check_golden(screen, name, args.out, args.golden)


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