#!/usr/bin/env python3
"""PPM loading, cropping and image comparison for UI flow region asserts.

Standalone (no VM, no other module in this repo needed): a screendump is
QEMU's own P6 PPM file, and everything below works straight off that file
plus a JSON reference. Pure Python plus Pillow only.

    imgtools.py compare REFS NAME SHOT.ppm
    imgtools.py capture REFS NAME SHOT.ppm X Y W H TOL   # TOL: exact|loose|structural
    imgtools.py hash SHOT.ppm
    imgtools.py bench [SHOT.ppm] [N]     # per-call phash/structural timing, N iterations
    imgtools.py selftest

`compare` and `capture` read/write a reference file: JSON mapping a region
name to {"x","y","w","h","tolerance","phash","structural"}. `capture` always
fills in both hashes so an existing entry's tolerance class can be changed
later without re-shooting the region.

Public API (imported directly by the driver, not just via the CLI above):

    load_ppm(path) -> (w, h, rgb_bytes)
    crop(img, x, y, w, h) -> img
    phash(img) -> int            8x8 DCT of a 32x32 grey downscale, 64-bit
    structural(img) -> bytes     16x16 edge/nonblack occupancy mask, 32 bytes
    compare(ref_entry, cur_img, tolerance) -> (ok, detail)

`img` is always the (w, h, rgb_bytes) triple load_ppm/crop return -- there is
no separate image type. rgb_bytes is tightly packed row-major RGB, 3 bytes
per pixel, no padding, matching both PPM P6 and PIL's "RGB" raw layout.
"""
import math
import os
import sys
import json

from PIL import Image, ImageFilter

# occupancy thresholds for structural(): a downsampled cell counts as
# "occupied" if it is meaningfully brighter than black, or sits on an edge.
# Both are 0..255 greyscale/edge-magnitude averages over the cell.
# structural(): a cell is occupied when its grey is this far from the crop's
# median (its background) or its FIND_EDGES energy exceeds the edge threshold.
# Validated on real captures: a switch knob left/right differs in 240/256
# cells, a dark card's icon and text stand out from its (7,13,29) ground.
DEVIATION_THRESHOLD = 28
EDGE_THRESHOLD = 24


# ---------------------------------------------------------------------------
# PPM
# ---------------------------------------------------------------------------

def load_ppm(path):
    """Read a binary PPM (P6), return (w, h, rgb_bytes).

    QEMU's screendump writes a comment-free header ("P6\\nW H\\n255\\n"
    then raw bytes) but the token reader here handles the general P6
    grammar -- whitespace runs and '#' comments -- since that costs
    nothing extra and means any other producer of P6 files just works.
    """
    with open(path, "rb") as f:
        data = f.read()

    def skip_ws_comments(p):
        while p < len(data):
            c = data[p]
            if c in b" \t\r\n":
                p += 1
            elif c == ord("#"):
                while p < len(data) and data[p] != ord("\n"):
                    p += 1
            else:
                break
        return p

    def read_token(p):
        p = skip_ws_comments(p)
        start = p
        while p < len(data) and data[p] not in b" \t\r\n":
            p += 1
        return data[start:p], p

    pos = 0
    magic, pos = read_token(pos)
    if magic != b"P6":
        raise ValueError(f"{path}: not a P6 PPM (magic={magic!r})")
    w_tok, pos = read_token(pos)
    h_tok, pos = read_token(pos)
    maxval_tok, pos = read_token(pos)
    w, h, maxval = int(w_tok), int(h_tok), int(maxval_tok)
    if maxval != 255:
        raise ValueError(f"{path}: unsupported PPM maxval {maxval} (only 255 handled)")
    # exactly one whitespace byte separates the header from the binary data
    pos += 1
    need = w * h * 3
    pixels = data[pos:pos + need]
    if len(pixels) != need:
        raise ValueError(f"{path}: truncated PPM, want {need} bytes got {len(pixels)}")
    return w, h, pixels


def crop(img, x, y, w, h):
    """Crop (W,H,rgb) to the x,y,w,h box, return a new (w,h,rgb) image."""
    width, height, data = img
    if x < 0 or y < 0 or w <= 0 or h <= 0 or x + w > width or y + h > height:
        raise ValueError(f"crop box {x},{y},{w}x{h} outside image {width}x{height}")
    row_bytes = w * 3
    out = bytearray(h * row_bytes)
    for row in range(h):
        src_off = ((y + row) * width + x) * 3
        dst_off = row * row_bytes
        out[dst_off:dst_off + row_bytes] = data[src_off:src_off + row_bytes]
    return w, h, bytes(out)


def _to_pil(img):
    w, h, data = img
    return Image.frombytes("RGB", (w, h), data)


# ---------------------------------------------------------------------------
# perceptual hash (phash): 8x8 DCT of a 32x32 grey downscale
# ---------------------------------------------------------------------------

def _dct_basis(n):
    """n x n DCT-II basis matrix. Unnormalised -- fine since phash only
    thresholds coefficients against each other, never compares magnitudes
    across images."""
    return [[math.cos(math.pi / n * (col + 0.5) * row) for col in range(n)] for row in range(n)]


_DCT32 = _dct_basis(32)

# phash() only ever reads the low-frequency _DCT_KEEP x _DCT_KEEP corner of
# the 32x32 DCT (see its docstring), so _dct2d_32 stops both passes at this
# many frequencies instead of computing all 1024 coefficients -- same 64
# coefficients out, roughly 6x less pure-Python multiply-add work per call.
_DCT_KEEP = 8


def _dct2d_32(rows):
    """2D DCT-II of a 32x32 list-of-lists via two separable 1D passes,
    truncated to the low-frequency _DCT_KEEP x _DCT_KEEP corner -- the only
    coefficients any caller reads."""
    n = 32
    # columns first: tmp[u][x] = DCT of column x at frequency u. Frequencies
    # u >= _DCT_KEEP never feed a returned coefficient, so skip them.
    tmp = [[0.0] * n for _ in range(_DCT_KEEP)]
    for u in range(_DCT_KEEP):
        b = _DCT32[u]
        for x in range(n):
            s = 0.0
            for y in range(n):
                s += rows[y][x] * b[y]
            tmp[u][x] = s
    # then rows: out[u][v] = DCT of tmp's row u at frequency v. Same cutoff
    # on v; the y/x summations stay full since each kept coefficient still
    # needs the whole 32-wide signal.
    out = [[0.0] * _DCT_KEEP for _ in range(_DCT_KEEP)]
    for u in range(_DCT_KEEP):
        row = tmp[u]
        for v in range(_DCT_KEEP):
            b = _DCT32[v]
            s = 0.0
            for x in range(n):
                s += row[x] * b[x]
            out[u][v] = s
    return out


def phash(img):
    """8x8-DCT perceptual hash of a 32x32 greyscale downscale, as a 64-bit
    int. Classic pHash recipe: downscale, DCT, keep the low-frequency 8x8
    corner, threshold each coefficient against the block's mean (DC term
    excluded from the mean -- it is just overall brightness and would bias
    every bit the same way)."""
    small = _to_pil(img).convert("L").resize((32, 32), Image.LANCZOS)
    px = small.load()
    rows = [[px[x, y] for x in range(32)] for y in range(32)]
    coeffs = _dct2d_32(rows)
    block = [coeffs[u][v] for u in range(_DCT_KEEP) for v in range(_DCT_KEEP)]
    mean = sum(block[1:]) / (len(block) - 1)
    bits = 0
    for c in block:
        bits = (bits << 1) | (1 if c > mean else 0)
    return bits


def hamming(a, b):
    """Bit-differences between two phash ints."""
    return bin(a ^ b).count("1")


# ---------------------------------------------------------------------------
# structural occupancy hash: 16x16 edge/nonblack mask
# ---------------------------------------------------------------------------

def structural(img):
    """16x16 binary occupancy mask as 32 bytes (256 bits, MSB first,
    row-major). A cell is "occupied" if it carries ink relative to the crop's
    OWN background -- its grey deviates from the crop's median by more than
    DEVIATION_THRESHOLD -- or sits on an edge, so the mask is robust to a
    recolour (still occupied) but sensitive to a shape disappearing (goes
    from occupied to empty).

    Relative to the median, not to black: the WardenOS page background is
    (7,13,29), grey 16, and an absolute non-black test read every cell of
    every region as occupied, so no structural check could ever fail
    (SDK #19, flare-edge #156). Measured on real captures with this rule: a
    Bluetooth switch off/on differs in 240 of 256 cells, a dark card reads
    its icon and text and nothing else."""
    grey = _to_pil(img).convert("L")
    edges = grey.filter(ImageFilter.FIND_EDGES)
    grey_small = grey.resize((16, 16), Image.BOX)
    edge_small = edges.resize((16, 16), Image.BOX)
    median = _median_grey(grey)
    gpx, epx = grey_small.load(), edge_small.load()
    bits = bytearray(32)
    idx = 0
    for y in range(16):
        for x in range(16):
            occupied = abs(gpx[x, y] - median) > DEVIATION_THRESHOLD or epx[x, y] > EDGE_THRESHOLD
            if occupied:
                bits[idx // 8] |= 1 << (7 - (idx % 8))
            idx += 1
    return bytes(bits)


def _median_grey(grey):
    """The crop's dominant luminance: the background of a card, the fill of a
    switch, whatever most of the pixels are."""
    hist = grey.histogram()
    total = sum(hist)
    acc = 0
    for value, count in enumerate(hist):
        acc += count
        if acc * 2 >= total:
            return value
    return 0


def _structural_diff(a, b):
    """Count of differing bits between two 32-byte occupancy masks."""
    return sum(bin(x ^ y).count("1") for x, y in zip(a, b))


# ---------------------------------------------------------------------------
# compare
# ---------------------------------------------------------------------------

def compare(ref_entry, cur_img, tolerance):
    """Compare cur_img (already cropped to the reference's box) against
    ref_entry under the given tolerance class. Returns (ok, detail); detail
    explains the verdict either way, for results.jsonl.

    exact:      phash hamming distance == 0
    loose:      phash hamming distance <= 6
    structural: occupancy mask bytes equal (no distance, no tolerance)
    """
    if tolerance in ("exact", "loose"):
        if "phash" not in ref_entry:
            return False, "reference missing 'phash' field"
        try:
            ref_hash = int(ref_entry["phash"], 16)
        except (TypeError, ValueError):
            return False, f"reference 'phash' is not valid hex: {ref_entry.get('phash')!r}"
        cur_hash = phash(cur_img)
        dist = hamming(ref_hash, cur_hash)
        limit = 0 if tolerance == "exact" else 6
        ok = dist <= limit
        detail = f"phash hamming={dist} limit={limit} ref={ref_hash:016x} cur={cur_hash:016x}"
        return ok, detail

    if tolerance == "structural":
        if "structural" not in ref_entry:
            return False, "reference missing 'structural' field"
        try:
            ref_bits = bytes.fromhex(ref_entry["structural"])
        except (TypeError, ValueError):
            return False, f"reference 'structural' is not valid hex: {ref_entry.get('structural')!r}"
        cur_bits = structural(cur_img)
        ok = ref_bits == cur_bits
        if ok:
            detail = "structural mask matches (0/256 cells differ)"
        else:
            diff = _structural_diff(ref_bits, cur_bits)
            detail = f"structural mask differs in {diff}/256 cells"
        return ok, detail

    return False, f"unknown tolerance class: {tolerance!r}"


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

def _load_refs(path):
    if not os.path.exists(path):
        return {}
    with open(path) as f:
        return json.load(f)


def _save_refs(path, refs):
    with open(path, "w") as f:
        json.dump(refs, f, indent=2, sort_keys=True)
        f.write("\n")


def cmd_compare(argv):
    if len(argv) != 3:
        print("usage: imgtools.py compare REFS NAME SHOT.ppm", file=sys.stderr)
        return 2
    refs_path, name, shot_path = argv
    refs = _load_refs(refs_path)
    if name not in refs:
        print(f"FAIL {name}: no such reference (run capture first)")
        return 1
    entry = refs[name]
    for key in ("x", "y", "w", "h", "tolerance"):
        if key not in entry:
            print(f"FAIL {name}: reference entry missing '{key}'")
            return 1
    w, h, data = load_ppm(shot_path)
    cropped = crop((w, h, data), entry["x"], entry["y"], entry["w"], entry["h"])
    ok, detail = compare(entry, cropped, entry["tolerance"])
    print(f"{'PASS' if ok else 'FAIL'} {name}: {detail}")
    return 0 if ok else 1


def cmd_capture(argv):
    if len(argv) != 8:
        print("usage: imgtools.py capture REFS NAME SHOT.ppm X Y W H TOL", file=sys.stderr)
        return 2
    refs_path, name, shot_path, x, y, w, h, tol = argv
    if tol not in ("exact", "loose", "structural"):
        print(f"usage: TOL must be exact|loose|structural, got {tol!r}", file=sys.stderr)
        return 2
    x, y, w, h = int(x), int(y), int(w), int(h)
    iw, ih, data = load_ppm(shot_path)
    cropped = crop((iw, ih, data), x, y, w, h)
    entry = {
        "x": x, "y": y, "w": w, "h": h,
        "tolerance": tol,
        "phash": f"{phash(cropped):016x}",
        "structural": structural(cropped).hex(),
    }
    refs = _load_refs(refs_path)
    refs[name] = entry
    _save_refs(refs_path, refs)
    print(f"captured {name}: {entry}")
    return 0


def cmd_hash(argv):
    if len(argv) != 1:
        print("usage: imgtools.py hash SHOT.ppm", file=sys.stderr)
        return 2
    w, h, data = load_ppm(argv[0])
    img = (w, h, data)
    print(f"{w}x{h} phash={phash(img):016x} structural={structural(img).hex()}")
    return 0


def cmd_bench(argv):
    """Time phash()/structural() -- the per-step cost of every
    assert_region/wait_region/capture_region in a flow run (wait_region polls
    at 0.5s intervals, calling back into these on every poll), so a change to
    the DCT size, the downscale filter, or the occupancy thresholds shows up
    as a number here instead of only as a slower flow run nobody investigates.

    usage: imgtools.py bench [SHOT.ppm] [N]

    SHOT.ppm: a real screendump to benchmark against a representative crop
    of; default is a synthetic 720x720 image (the virt.fragment screendump
    size) so bench never depends on a committed fixture or a live rig.
    N: iterations per function, default 200.
    """
    import timeit

    if len(argv) > 2:
        print("usage: imgtools.py bench [SHOT.ppm] [N]", file=sys.stderr)
        return 2
    shot_path = next((a for a in argv if not a.isdigit()), None)
    n = int(next((a for a in argv if a.isdigit()), "200"))

    if shot_path:
        w, h, data = load_ppm(shot_path)
    else:
        w, h = 720, 720
        data = _make_test_rgb(w, h, (220, 20, 20))
    img = (w, h, data)
    # a representative crop, not the whole screen: every real region assert
    # crops first, and structural()'s edge filter cost scales with crop
    # size even though phash's fixed 32x32 downscale mostly doesn't.
    cw, ch = min(200, w), min(200, h)
    region = crop(img, 0, 0, cw, ch)

    def percentiles(samples):
        s = sorted(samples)
        p50 = s[min(len(s) - 1, int(len(s) * 0.50))]
        p95 = s[min(len(s) - 1, int(len(s) * 0.95))]
        return s[0], p50, p95, s[-1]

    print(f"bench: {cw}x{ch} region, n={n} iterations")
    for label, fn in (("phash", lambda: phash(region)),
                       ("structural", lambda: structural(region))):
        samples = timeit.repeat(fn, repeat=n, number=1)
        lo, p50, p95, hi = percentiles(samples)
        print(f"  {label:<10} min={lo * 1000:7.3f}ms  p50={p50 * 1000:7.3f}ms  "
              f"p95={p95 * 1000:7.3f}ms  max={hi * 1000:7.3f}ms")
    return 0


# ---------------------------------------------------------------------------
# selftest
# ---------------------------------------------------------------------------

def _make_test_rgb(w, h, colour, shift=(0, 0)):
    """Black canvas with a coloured square near one corner: the self-test's
    synthetic images, and bench's default when it isn't given a real
    SHOT.ppm. colour=None means no square at all (removed shape)."""
    from PIL import ImageDraw
    img = Image.new("RGB", (w, h), (0, 0, 0))
    if colour is not None:
        draw = ImageDraw.Draw(img)
        ox, oy = shift
        x0, y0 = 8 + ox, 8 + oy
        draw.rectangle([x0, y0, x0 + 20, y0 + 20], fill=colour)
    return img.tobytes()


def _write_ppm(path, w, h, rgb_bytes):
    with open(path, "wb") as f:
        f.write(f"P6\n{w} {h}\n255\n".encode("ascii"))
        f.write(rgb_bytes)


def _dct2d_32_brute(rows):
    """Unoptimised reference DCT: the full 32x32 transform with no early
    cutoff. selftest()-only, so _dct2d_32's _DCT_KEEP truncation has
    something independent to be checked against -- a future edit that moves
    the cutoff on the wrong loop would otherwise change captured phash bits
    with nothing catching it."""
    n = 32
    tmp = [[0.0] * n for _ in range(n)]
    for u in range(n):
        b = _DCT32[u]
        for x in range(n):
            s = 0.0
            for y in range(n):
                s += rows[y][x] * b[y]
            tmp[u][x] = s
    out = [[0.0] * n for _ in range(n)]
    for u in range(n):
        row = tmp[u]
        for v in range(n):
            b = _DCT32[v]
            s = 0.0
            for x in range(n):
                s += row[x] * b[x]
            out[u][v] = s
    return out


def selftest():
    import tempfile

    w, h = 128, 128
    failures = []

    def check(label, cond):
        print(f"[{'ok' if cond else 'FAIL'}] {label}")
        if not cond:
            failures.append(label)

    with tempfile.TemporaryDirectory() as tmp:
        refs_path = os.path.join(tmp, "refs.json")

        base_path = os.path.join(tmp, "base.ppm")
        _write_ppm(base_path, w, h, _make_test_rgb(w, h, (220, 20, 20)))

        # capture the same corner box under all three tolerance classes
        for tol in ("exact", "loose", "structural"):
            rc = main(["capture", refs_path, f"corner-{tol}", base_path, "4", "4", "32", "32", tol])
            check(f"capture corner-{tol} succeeds", rc == 0)

        # load_ppm/crop sanity: the drawn pixel lands where expected
        bw, bh, bdata = load_ppm(base_path)
        check("load_ppm reads back the written size", (bw, bh) == (w, h))
        cropped = crop((bw, bh, bdata), 4, 4, 32, 32)
        px_off = ((8 - 4) * 32 + (8 - 4)) * 3
        check("crop keeps the drawn pixel in place",
              cropped[2][px_off:px_off + 3] == bytes((220, 20, 20)))

        # _dct2d_32's _DCT_KEEP truncation must land on the exact same
        # coefficients an untruncated 32x32 DCT would produce
        small = _to_pil((bw, bh, bdata)).convert("L").resize((32, 32), Image.LANCZOS)
        px = small.load()
        dct_rows = [[px[x, y] for x in range(32)] for y in range(32)]
        fast, brute = _dct2d_32(dct_rows), _dct2d_32_brute(dct_rows)
        check("_dct2d_32's truncated corner matches the untruncated DCT",
              all(fast[u][v] == brute[u][v]
                  for u in range(_DCT_KEEP) for v in range(_DCT_KEEP)))

        # identical shot: exact must pass
        same_path = os.path.join(tmp, "same.ppm")
        _write_ppm(same_path, w, h, _make_test_rgb(w, h, (220, 20, 20)))
        rc = main(["compare", refs_path, "corner-exact", same_path])
        check("exact passes for an identical shot", rc == 0)

        # small shift: exact must fail, loose must still pass
        shift_path = os.path.join(tmp, "shift.ppm")
        _write_ppm(shift_path, w, h, _make_test_rgb(w, h, (220, 20, 20), shift=(1, 1)))
        rc = main(["compare", refs_path, "corner-exact", shift_path])
        check("exact fails for a slight shift", rc == 1)
        rc = main(["compare", refs_path, "corner-loose", shift_path])
        check("loose passes for a slight shift", rc == 0)

        # colour change, same footprint: structural must still pass
        colour_path = os.path.join(tmp, "colour.ppm")
        _write_ppm(colour_path, w, h, _make_test_rgb(w, h, (20, 20, 220)))
        rc = main(["compare", refs_path, "corner-structural", colour_path])
        check("structural passes for a colour change", rc == 0)

        # shape removed entirely: structural must fail
        removed_path = os.path.join(tmp, "removed.ppm")
        _write_ppm(removed_path, w, h, _make_test_rgb(w, h, None))
        rc = main(["compare", refs_path, "corner-structural", removed_path])
        check("structural fails when the shape is removed", rc == 1)

        # bad-path coverage
        rc = main(["compare", refs_path, "no-such-name", base_path])
        check("compare on an unknown reference name fails loudly", rc == 1)
        rc = main(["hash", base_path])
        check("hash on a real ppm succeeds", rc == 0)
        rc = main(["bench", base_path, "3"])
        check("bench runs to completion", rc == 0)

    if failures:
        print(f"\n{len(failures)} check(s) failed:")
        for label in failures:
            print(f"  - {label}")
        return 1
    print("\nall checks passed")
    return 0


# ---------------------------------------------------------------------------

def main(argv):
    if not argv:
        print(__doc__, file=sys.stderr)
        return 2
    cmd, rest = argv[0], argv[1:]
    if cmd == "selftest":
        return selftest()
    handlers = {"compare": cmd_compare, "capture": cmd_capture, "hash": cmd_hash, "bench": cmd_bench}
    if cmd not in handlers:
        print(f"unknown command: {cmd!r}", file=sys.stderr)
        return 2
    # A missing file, unreadable ref JSON, or malformed PPM/crop box is a
    # prerequisite failure -- report it plainly and fail loudly (exit 1),
    # never let it fall through as a silent pass or a raw traceback.
    try:
        return handlers[cmd](rest)
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        print(f"FATAL: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))