#!/usr/bin/env python3
"""keysRow.py - read KEYS's verification row off a screenshot.

KEYS (examples/keys) draws every character it received on its row 198 as
pixel nibbles: the sentinel A 5 A 5, the count byte, a zero byte, then each
character, two pixels a byte, high nibble first. This finds that row in a
screenshot (on a desktop the window is wherever DVX put it) and checks it
holds exactly TARGET.

    keysRow.py PNG TARGET [--size WxH] [--at X,Y]

--size: the screenshot must be that size (the full-screen mode, not the
desktop). --at: the row must start exactly there (the stage's own row 198
at 1x on a full screen). Prints the row it found; exit 0 when it is right.
"""

import sys

from PIL import Image

# KEYS palette 0 (buildPalette): a gray ramp with four demo colours and red
# at 15, 12-bit $0RGB, each channel widened x17 as the DVX port's tables do.
PAL12 = [(i << 8) | (i << 4) | i for i in range(16)]
PAL12[0], PAL12[1], PAL12[2], PAL12[3] = 0x000, 0x333, 0x0F0, 0xFFF
PAL12[15] = 0xF00
PAL_RGB = [(((v >> 8) & 0xF) * 17, ((v >> 4) & 0xF) * 17, (v & 0xF) * 17) for v in PAL12]
# A 16-bit desktop (5-6-5) is off by at most 8 a channel after the round
# trip, the 6-bit VGA DAC by at most 3, while the palette entries sit 17
# apart: nearest is exact.
MAX_DIST = 3 * 9 ** 2
SENTINEL = (0xA, 0x5, 0xA, 0x5)


def main(argv):
    if len(argv) < 3:
        sys.stderr.write(__doc__)
        return 2
    png, target = argv[1], argv[2]
    size = None
    at = None
    args = argv[3:]
    while args:
        if args[0] == "--size" and len(args) > 1:
            size = tuple(int(n) for n in args[1].split("x"))
        elif args[0] == "--at" and len(args) > 1:
            at = tuple(int(n) for n in args[1].split(","))
        else:
            sys.stderr.write(__doc__)
            return 2
        args = args[2:]

    img = Image.open(png).convert("RGB")
    w, h = img.size
    px = img.load()
    if size is not None and img.size != size:
        print("keysRow: FAIL - %s is %dx%d, want %dx%d" % (png, w, h, size[0], size[1]))
        return 1

    nearest = {}

    def nib(x, y):
        rgb = px[x, y]
        if rgb not in nearest:
            best, bestDist = 0, 1 << 30
            for n, pal in enumerate(PAL_RGB):
                d = sum((c - p) ** 2 for c, p in zip(rgb, pal))
                if d < bestDist:
                    best, bestDist = n, d
            nearest[rgb] = best if bestDist <= MAX_DIST else -1
        return nearest[rgb]

    expect = bytes([0xA5, 0xA5, len(target), 0x00]) + target.encode("ascii")
    found = None
    for y in range(h):
        for x in range(w - 2 * len(expect)):
            if tuple(nib(x + i, y) for i in range(len(SENTINEL))) == SENTINEL:
                row = bytes(((nib(x + 2 * i, y) & 0xF) << 4) | (nib(x + 2 * i + 1, y) & 0xF) for i in range(len(expect)))
                if row[:4] == expect[:4]:
                    found = (x, y, row)
                    break
        if found:
            break
    if not found:
        print("keysRow: FAIL - no verification row (A5A5 %02X00) on the screen" % len(target))
        return 1
    x, y, row = found
    print("keysRow: row at (%d,%d)=%s" % (x, y, row.hex().upper()))
    if row != expect:
        print("keysRow: FAIL - screen row %s, want %s" % (row.hex().upper(), expect.hex().upper()))
        return 1
    if at is not None and (x, y) != at:
        print("keysRow: FAIL - the row is at (%d,%d), want (%d,%d)" % (x, y, at[0], at[1]))
        return 1
    return 0


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