#!/usr/bin/env python3
# Author:  cirakg & claude ai
# Summary: Drive the app in the headless simulator with scripted input and check a screenshot of every screen.
"""Scripted UI flow for the UI-test build (README 17.4).

Usage (inside the dev container, with the simulator running and the
UI-test build launched by docker/scripts/screenshots.sh):
    ui_flow.py --device fr70 --device-dir DIR --log APP.log --window ID
               --out build/screenshots/fr70 --golden test/golden/fr70 [--approve]

Input goes to the simulator window with xdotool: clicks on the device
drawing's buttons (positions from DIR/simulator.json) and taps and swipes on
the screen. Watches without a button for an action use touch instead (Venu
Sq 2 has no UP/DOWN and no select button).

The app (UI-test build, source/util/TestMode.mc) starts with fake entries,
the PIN 123456 and a frozen clock, and prints "DIAG shown <screen>" after
drawing each screen, so every step waits for the screen it expects instead
of sleeping. Menus are native Menu2 views that print nothing; for those the
script waits a fixed time.

The flow covers: the results of an Aegis import sent through the settings
form (the fixture sets the form, as the phone would), the list, a TOTP entry, a HOTP entry and its next code, the
PIN screen, a wrong PIN, a correct PIN opening a locked entry, auto-lock
closing it, the menu, the About screen (watches with buttons; touch-only
menus have no focus to move), the lockout after 5 wrong PINs, and an EAN-13
code whose screenshot must decode with zxing-cpp (README 12.4).
Each screenshot is compared with test/golden/<device>/<screen>.png like
sim_screenshot.py does (or saved there with --approve).

Exit: 0 all screens match, 1 a screen differs or is missing a golden image,
2 the flow failed (a screen didn't appear in time).
"""

from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path

import sim_screenshot as shots
import zxingcpp

EXIT_OK, EXIT_DIFFERENT, EXIT_FLOW = 0, 1, 2
SCREEN_TIMEOUT_S = 20.0
AUTO_LOCK_TIMEOUT_S = 25.0  # the UI-test build locks after 10 s without input
IMPORT_TIMEOUT_S = 90.0  # the 100-entry fixture import runs on a timer after start-up
IMPORT_MEMORY_BUDGET = 0.70  # README 18: peak memory at most 70 % of the watch-app memory
MENU_SETTLE_S = 1.5  # Menu2 has no "shown" line
SETTLE_S = 0.4  # after a "shown" line, before the screenshot
CLICK_PAUSE_S = 0.25
HOLD_S = 1.6  # long enough for a hold (menu) on every device
FIXTURE_PIN = "123456"  # source/app/Fixtures.mc
WRONG_PIN = "000000"
FAILURES_UNTIL_WAIT = 5  # source/storage/Lockout.mc: the 5th failure starts a delay
SELECTOR_OK = 10  # source/ui/PinView.mc: selector positions 0-9, then OK
GYM_CARD = "5901234123457"  # source/app/Fixtures.mc: EAN-13 with its check digit
MEMBER_QR = (  # source/app/Fixtures.mc: a version 6 QR
    "https://example.com/member/fake-0042?name=Alice%20Example&card=0000111122223333&valid=2027-12-31&level=gold"
)
PARCEL_DM = "PARCEL-FAKE-0042-7731"  # source/app/Fixtures.mc: Data Matrix
TICKET_PDF417 = "Boarding pass FAKE 0042 seat 12A"  # source/app/Fixtures.mc: PDF417
QR_TIMEOUT_S = 30.0  # the QR is encoded in timer steps the first time


class FlowError(Exception):
    """The app didn't reach an expected screen."""


@dataclass(frozen=True)
class Key:
    x: int
    y: int
    behavior: str
    hold: bool


def load_keys(device_dir: Path) -> dict[str, Key]:
    """Button centres in device-drawing coordinates, by button id (enter, up, down, esc, menu)."""
    sim = json.loads((device_dir / "simulator.json").read_text(encoding="utf-8"))
    keys: dict[str, Key] = {}
    for key in sim.get("keys", []):
        loc = key["location"]
        keys[str(key["id"])] = Key(
            int(loc["x"]) + int(loc["width"]) // 2,
            int(loc["y"]) + int(loc["height"]) // 2,
            str(key.get("behavior", "")),
            bool(key.get("isHold", False)),
        )
    return keys


class AppLog:
    """The app's console output, read as it grows."""

    def __init__(self, path: Path) -> None:
        self.path = path
        self.mark = 0

    def lines(self) -> list[str]:
        try:
            return self.path.read_text(encoding="utf-8", errors="replace").splitlines()
        except OSError:
            return []

    def set_mark(self) -> None:
        """Later waits only look at lines written after this point."""
        self.mark = len(self.lines())

    def wait_for(self, pattern: str, timeout: float) -> re.Match[str]:
        regex = re.compile(pattern)
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            for line in self.lines()[self.mark :]:
                match = regex.search(line)
                if match:
                    return match
            time.sleep(0.05)
        raise FlowError(f"timed out after {timeout:.0f} s waiting for /{pattern}/")

    def latest(self, pattern: str) -> re.Match[str] | None:
        regex = re.compile(pattern)
        found = None
        for line in self.lines():
            match = regex.search(line)
            if match:
                found = match
        return found


class Simulator:
    """Input to and screenshots of the simulator window."""

    def __init__(self, window: str, device_dir: Path, raw_dir: Path) -> None:
        self.window = window
        self.device_dir = device_dir
        self.display = shots.load_display(device_dir)
        self.keys = load_keys(device_dir)
        self.raw_dir = raw_dir
        self.offset = self._find_offset()

    def capture(self, name: str) -> Path:
        path = self.raw_dir / f"{name}-window.png"
        path.parent.mkdir(parents=True, exist_ok=True)
        subprocess.run(["import", "-window", self.window, str(path)], check=True)
        return path

    def _find_offset(self) -> tuple[int, int]:
        raw = shots.open_image(self.capture("offset"))
        device = shots.open_image(self.device_dir / self.display.image)
        dx, dy, score = shots.find_offset(raw, shots.bezel_only(device, self.display))
        if score > shots.MAX_BEZEL_DIFF:
            raise FlowError(f"device drawing not found in the simulator window (difference {score:.1f})")
        return dx, dy

    def _xdotool(self, *args: str) -> None:
        subprocess.run(["xdotool", *args], check=True)

    def _move(self, wx: int, wy: int) -> None:
        self._xdotool("mousemove", "--window", self.window, str(wx), str(wy))

    def _to_window(self, dx: int, dy: int) -> tuple[int, int]:
        """Device-drawing coordinates to window coordinates."""
        return self.offset[0] + dx, self.offset[1] + dy

    def has(self, key_id: str, behavior: str | None = None) -> bool:
        key = self.keys.get(key_id)
        return key is not None and (behavior is None or key.behavior == behavior)

    def press(self, key_id: str, hold: bool = False) -> None:
        key = self.keys[key_id]
        self._move(*self._to_window(key.x, key.y))
        if hold:
            self._xdotool("mousedown", "1")
            time.sleep(HOLD_S)
            self._xdotool("mouseup", "1")
        else:
            self._xdotool("click", "1")
        time.sleep(CLICK_PAUSE_S)

    def tap(self, sx: int, sy: int) -> None:
        """Tap at screen coordinates (0, 0 = top left of the watch screen)."""
        self._move(*self._to_window(self.display.x + sx, self.display.y + sy))
        self._xdotool("click", "1")
        time.sleep(CLICK_PAUSE_S)

    def long_press(self, sx: int, sy: int) -> None:
        self._move(*self._to_window(self.display.x + sx, self.display.y + sy))
        self._xdotool("mousedown", "1")
        time.sleep(HOLD_S)
        self._xdotool("mouseup", "1")
        time.sleep(CLICK_PAUSE_S)

    def swipe(self, sx0: int, sy0: int, sx1: int, sy1: int) -> None:
        self._move(*self._to_window(self.display.x + sx0, self.display.y + sy0))
        self._xdotool("mousedown", "1")
        steps = 8
        for i in range(1, steps + 1):
            x = sx0 + (sx1 - sx0) * i // steps
            y = sy0 + (sy1 - sy0) * i // steps
            self._move(*self._to_window(self.display.x + x, self.display.y + y))
            time.sleep(0.02)
        self._xdotool("mouseup", "1")
        time.sleep(CLICK_PAUSE_S * 2)


class Flow:
    """The abstract actions of README 13.4, mapped to this device's buttons or touch screen."""

    def __init__(self, sim: Simulator, log: AppLog, out: Path, golden: Path, approve: bool) -> None:
        self.sim = sim
        self.log = log
        self.out = out
        self.golden = golden
        self.approve = approve
        self.results: dict[str, int] = {}
        self.selector = 0  # the PIN digit selector's position (button-only watches)
        touch = log.wait_for(r"DIAG device touch=(\d)", SCREEN_TIMEOUT_S)
        self.touch = touch.group(1) == "1"

    # -- navigation ---------------------------------------------------------

    def _w(self) -> int:
        return self.sim.display.width

    def _h(self) -> int:
        return self.sim.display.height

    def next(self) -> None:
        if self.sim.has("down"):
            self.sim.press("down")
        else:
            self.sim.swipe(self._w() // 2, self._h() * 3 // 4, self._w() // 2, self._h() // 4)

    def prev(self) -> None:
        if self.sim.has("up"):
            self.sim.press("up")
        else:
            self.sim.swipe(self._w() // 2, self._h() // 4, self._w() // 2, self._h() * 3 // 4)

    def back(self) -> None:
        self.sim.press("esc")

    def menu(self) -> None:
        """Hold MENU, or long-press the screen on watches without one (Venu X1)."""
        if self.sim.has("menu"):
            self.sim.press("menu", hold=True)
        else:
            self.sim.long_press(self._w() // 2, self._h() // 2)

    def open_focused_row(self) -> None:
        if self.sim.has("enter", "onSelect"):
            self.sim.press("enter")
            return
        focus = self.log.latest(r"DIAG list_focus y=(\d+)")
        if focus is None:
            raise FlowError("the list didn't report its focused row")
        self.sim.tap(self._w() // 2, int(focus.group(1)))

    def select(self) -> None:
        """START, or a tap in the middle of the screen."""
        if self.sim.has("enter", "onSelect"):
            self.sim.press("enter")
        else:
            self.sim.tap(self._w() // 2, self._h() // 2)

    def enter_pin(self, pin: str) -> None:
        if self.touch:
            keys = self.log.latest(r"DIAG pin_keys (\S+)")
            if keys is None:
                raise FlowError("the PIN screen didn't report its keypad")
            centres = {}
            for item in keys.group(1).split(";"):
                value, xy = item.split("=")
                x, y = xy.split(",")
                centres[int(value)] = (int(x), int(y))
            for digit in pin:
                self.sim.tap(*centres[int(digit)])
            self.sim.tap(*centres[-2])  # OK
            return
        for position in [int(d) for d in pin] + [SELECTOR_OK]:
            steps = (position - self.selector) % (SELECTOR_OK + 1)
            for _ in range(steps):
                self.sim.press("down")
            self.selector = position
            self.sim.press("enter")

    # -- checking -----------------------------------------------------------

    def expect(self, screen: str, timeout: float = SCREEN_TIMEOUT_S) -> None:
        self.log.wait_for(rf"DIAG shown {re.escape(screen)}\s*$", timeout)
        time.sleep(SETTLE_S)

    def opened_pin_screen(self) -> None:
        """A new PIN screen starts with its selector at 0."""
        self.selector = 0

    def act(self) -> None:
        """Call before an input whose result `expect` should wait for."""
        self.log.set_mark()

    def shot(self, name: str, expect_code: str | None = None) -> None:
        raw = self.sim.capture(name)
        out = self.out / f"{name}.png"
        label = f"{self.sim.device_dir.name}/{name}"
        try:
            screen = shots.extract_screen(raw, self.sim.device_dir, out)
        except shots.ScreenshotError as err:
            print(f"[{label}] {err}")
            self.results[name] = EXIT_DIFFERENT
            return
        if expect_code is not None:
            # README 12.4 check 3: the code on the screenshot must decode.
            read = [r.text for r in zxingcpp.read_barcodes(screen)]
            if expect_code not in read:
                print(f"[{label}] the barcode didn't decode to {expect_code!r}: read {read}")
                self.results[name] = EXIT_DIFFERENT
                return
            print(f"[{label}] barcode decodes to {expect_code!r}")
        golden = self.golden / f"{name}.png"
        if self.approve:
            shots.save_png(screen, golden, f"Approved simulator screen '{name}' of {self.sim.device_dir.name}")
            print(f"[{label}] approved as {golden}")
            self.results[name] = EXIT_OK
            return
        self.results[name] = shots.check_golden(screen, label, out, golden)


def flow_import(f: Flow) -> None:
    """The fixture's 100-entry Aegis import through the settings form opens its
    results; its peak memory must stay within the budget."""
    f.expect("results", IMPORT_TIMEOUT_S)
    stats = f.log.latest(r"DIAG import entries=(\d+) wall_ms=(\d+) mem_peak=(\d+) mem_total=(\d+)")
    if stats is None:
        raise FlowError("the import didn't report its statistics")
    entries, wall_ms, peak, total = (int(v) for v in stats.groups())
    fraction = peak / total
    print(f"[{f.sim.device_dir.name}] import of {entries} entries: {wall_ms} ms, peak memory {fraction:.1%}")
    if fraction > IMPORT_MEMORY_BUDGET:
        raise FlowError(f"import peak memory {fraction:.1%} is over {IMPORT_MEMORY_BUDGET:.0%}")
    f.shot("import_results")
    f.act()
    f.back()
    f.expect("list")


def flow_entries(f: Flow) -> None:
    """The list, a TOTP entry, a HOTP entry and its next code."""
    f.shot("list")
    f.act()
    f.open_focused_row()  # GitHub, TOTP
    f.expect("otp")
    f.shot("otp")
    f.act()
    f.back()
    f.expect("list")
    for _ in range(3):  # Google, Bank, VPN
        f.next()
    f.act()
    f.open_focused_row()  # VPN, HOTP
    f.expect("otp_hotp")
    f.shot("otp_hotp")
    f.act()
    f.select()  # next HOTP code
    f.expect("otp_hotp")
    f.shot("otp_hotp_next")
    f.act()
    f.back()
    f.expect("list")


def flow_locked_entry(f: Flow) -> None:
    """Bank is locked: PIN screen, a wrong PIN, the right PIN, then auto-lock."""
    f.prev()  # from VPN back to Bank
    f.act()
    f.open_focused_row()
    f.expect("pin")
    f.opened_pin_screen()
    f.shot("pin")
    f.act()
    f.enter_pin(WRONG_PIN)
    f.expect("pin_error")
    f.shot("pin_error")
    f.act()
    f.enter_pin(FIXTURE_PIN)
    f.expect("otp")
    f.shot("otp_locked")
    # Auto-lock closes the locked entry.
    f.act()
    f.expect("list", AUTO_LOCK_TIMEOUT_S)


def flow_menu(f: Flow) -> None:
    """The menu, and About on watches whose menu has a focus to move."""
    f.menu()
    time.sleep(MENU_SETTLE_S)
    f.shot("menu")
    f.act()
    if f.sim.has("down"):
        # Locked with a PIN: Change PIN, Reset PIN, Settings results, Entry
        # IDs, Clock status, About.
        for _ in range(5):
            f.next()
        time.sleep(MENU_SETTLE_S)
        f.act()
        f.select()
        f.expect("about")
        f.shot("about")
        f.act()
    # Touch-only watches show the menu as a freely scrolling list without
    # focus, so there is no reliable "select About"; just close it.
    f.back()
    f.expect("list")


def flow_lockout(f: Flow) -> None:
    """Five wrong PINs start the 30 s wait."""
    f.act()
    f.open_focused_row()
    f.expect("pin")
    f.opened_pin_screen()
    for attempt in range(1, FAILURES_UNTIL_WAIT + 1):
        f.act()
        f.enter_pin(WRONG_PIN)
        f.expect("pin_wait" if attempt == FAILURES_UNTIL_WAIT else "pin_error")
    f.shot("pin_wait")


def flow_code(f: Flow) -> None:
    """The EAN-13 gym card (index 6), the QR (7), Data Matrix (8) and PDF417 (9)
    cards: their screenshots must decode."""
    f.act()
    f.back()  # leave the waiting PIN screen
    f.expect("list")
    for _ in range(4):  # Bank (2) to the gym card (6)
        f.next()
    f.act()
    f.open_focused_row()
    f.expect("code")
    f.shot("code", expect_code=GYM_CARD)
    # The QR card (index 7) is encoded on the watch as a job, then shown.
    f.act()
    f.back()
    f.expect("list")
    f.next()
    f.act()
    f.open_focused_row()
    f.expect("code", QR_TIMEOUT_S)
    encoded = f.log.latest(r"DIAG matrix_encoded type=\d+ size=(\S+) ms=(\d+)")
    if encoded is not None:
        print(f"[{f.sim.device_dir.name}] QR {encoded.group(1)} encoded in {encoded.group(2)} ms")
    f.shot("code_qr", expect_code=MEMBER_QR)
    # Data Matrix (8) and PDF417 (9), also encoded as jobs the first time.
    for name, expected in (("code_datamatrix", PARCEL_DM), ("code_pdf417", TICKET_PDF417)):
        f.act()
        f.back()
        f.expect("list")
        f.next()
        f.act()
        f.open_focused_row()
        f.expect("code", QR_TIMEOUT_S)
        f.shot(name, expect_code=expected)


def run_flow(f: Flow) -> None:
    """The scripted session; see the module docstring."""
    flow_import(f)
    flow_entries(f)
    flow_locked_entry(f)
    flow_menu(f)
    flow_lockout(f)
    flow_code(f)


def find_key(data: object, key: str) -> object:
    """The first value stored under `key` anywhere in a JSON structure, or None."""
    if isinstance(data, dict):
        if key in data:
            return data[key]
        data = list(data.values())
    if isinstance(data, list):
        for item in data:
            found = find_key(item, key)
            if found is not None:
                return found
    return None


def open_from_glance(window: str, device_dir: Path, raw_dir: Path) -> None:
    """Open the full app from its glance, as found in the simulator: a tap on
    it on touch screens; on button-only watches DOWN (which selects the
    glance) and then START. START alone does nothing in either case."""
    sim = Simulator(window, device_dir, raw_dir)
    config = json.loads((device_dir / "simulator.json").read_text(encoding="utf-8"))
    area = find_key(config.get("glance", {}), "contentArea")
    if find_key(config, "isTouch") is True and isinstance(area, dict):
        sim.tap(int(area["x"]) + int(area["width"]) // 2, int(area["y"]) + int(area["height"]) // 2)
    else:
        sim.press("down")
        time.sleep(1.0)
        sim.press("enter")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Scripted UI flow with screenshot checks.")
    parser.add_argument("--device-dir", type=Path, required=True)
    parser.add_argument("--window", required=True, help="X window id of the simulator")
    parser.add_argument("--out", type=Path, required=True, help="directory for the screenshots")
    parser.add_argument("--log", type=Path, help="the app's console log (monkeydo output)")
    parser.add_argument("--golden", type=Path, help="directory of the approved images")
    parser.add_argument("--approve", action="store_true")
    parser.add_argument("--open-from-glance", action="store_true", help="only open the app from its glance")
    args = parser.parse_args(argv)
    name = args.device_dir.name
    if args.open_from_glance:
        try:
            open_from_glance(args.window, args.device_dir, args.out / "raw")
        except (FlowError, subprocess.CalledProcessError, OSError) as err:
            print(f"[{name}] could not open the app from its glance: {err}")
            return EXIT_FLOW
        return EXIT_OK
    if args.log is None or args.golden is None:
        parser.error("--log and --golden are required for the flow")
    try:
        sim = Simulator(args.window, args.device_dir, args.out / "raw")
        flow = Flow(sim, AppLog(args.log), args.out, args.golden, args.approve)
        run_flow(flow)
    except (FlowError, subprocess.CalledProcessError, OSError) as err:
        print(f"[{name}] UI flow failed: {err}")
        return EXIT_FLOW
    failed = sorted(screen for screen, code in flow.results.items() if code != EXIT_OK)
    if failed:
        print(f"[{name}] {len(failed)} of {len(flow.results)} screens need attention: {', '.join(failed)}")
        return EXIT_DIFFERENT
    print(f"[{name}] all {len(flow.results)} screens match")
    return EXIT_OK


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