#!/usr/bin/env python3
# Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
# See LICENSE.md for license information.
# SPDX-License-Identifier: MIT

"""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, at the window's
scale: every stage pixel a SCALE x SCALE block) and checks it holds exactly
TARGET, and that every block of it is whole: each of its screen pixels the
same colour.

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

--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). --scale: the row must be at that scale (else the
first of 1, 2, 3 it is found at). 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 15- or 16-bit desktop 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. (Every KEYS colour is also one of DVX's 256-colour
# desktop palette, its cube or its grey ramp, so 8 bits is exact too.)
MAX_DIST = 3 * 9 ** 2
NO_NIBBLE = -1
SENTINEL = (0xA, 0x5, 0xA, 0x5)
SCALES = (1, 2, 3)


def nibbleMap(img):
    """Every pixel's KEYS palette index, NO_NIBBLE where no entry is near."""
    nearest = {}
    for rgb in set(img.getdata()):
        best, bestDist = NO_NIBBLE, MAX_DIST + 1
        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
    return [nearest[rgb] for rgb in img.getdata()]


def findRow(nibs, pixels, w, h, scale, expect):
    """The (x, y, bytes) of the first row at scale whose header matches."""
    header = expect[:4]
    span = 2 * len(expect) * scale
    for y in range(h):
        base = y * w
        for x in range(w - span + 1):
            if nibs[base + x] != SENTINEL[0]:
                continue
            if tuple(nibs[base + x + i * scale] for i in range(len(SENTINEL))) != SENTINEL:
                continue
            row = bytes(((nibs[base + x + 2 * i * scale] & 0xF) << 4) | (nibs[base + x + (2 * i + 1) * scale] & 0xF) for i in range(len(expect)))
            if row[:4] == header:
                return x, y, row
    return None


def splitBlocks(pixels, w, h, x, y, scale, count):
    """How many of the count stage pixels from (x, y) are not whole blocks."""
    split = 0
    for i in range(count):
        left = x + i * scale
        want = pixels[y * w + left]
        block = [pixels[(y + dy) * w + left + dx] for dy in range(scale) if y + dy < h for dx in range(scale)]
        if len(block) != scale * scale or any(p != want for p in block):
            split += 1
    return split


def main(argv):
    if len(argv) < 3:
        sys.stderr.write(__doc__)
        return 2
    png, target = argv[1], argv[2]
    size = None
    at = None
    scales = SCALES
    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(","))
        elif args[0] == "--scale" and len(args) > 1:
            scales = (int(args[1]),)
        else:
            sys.stderr.write(__doc__)
            return 2
        args = args[2:]

    img = Image.open(png).convert("RGB")
    w, h = img.size
    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

    pixels = list(img.getdata())
    nibs = nibbleMap(img)
    expect = bytes([0xA5, 0xA5, len(target), 0x00]) + target.encode("ascii")
    found = None
    for scale in scales:
        found = findRow(nibs, pixels, w, h, scale, expect)
        if found:
            break
    if not found:
        print("keysRow: FAIL - no verification row (A5A5 %02X00) on the screen at %s" % (len(target), " or ".join("%dx" % s for s in scales)))
        return 1
    x, y, row = found
    print("keysRow: row at (%d,%d) at %dx=%s" % (x, y, scale, row.hex().upper()))
    if row != expect:
        print("keysRow: FAIL - screen row %s, want %s" % (row.hex().upper(), expect.hex().upper()))
        return 1
    split = splitBlocks(pixels, w, h, x, y, scale, 2 * len(expect))
    if split:
        print("keysRow: FAIL - %d of the row's %d stage pixels are not whole %dx%d blocks" % (split, 2 * len(expect), scale, scale))
        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))