#!/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

"""joyScreen.py - read both joysticks' state off a screenshot of JOY.

JOY (examples/joy) shows each stick in a frame: a dot that sits in the
middle at rest and moves DOT_TRAVEL stage pixels per full axis
(jlJoystickX/Y * DOT_TRAVEL / 127, grey at rest, green away from it), two
button squares (yellow while down) and a patch under them, green when the
port reports a joystick and red when it does not. This finds the stage in
a screenshot (on a desktop the window is wherever DVX put it, at the
window's scale) by the two patches, and prints each stick's state as five
characters: connected (1/0), X and Y (- 0 +: the dot at least half its
travel left or up, in the middle, at least half its travel right or
down), buttons 0 and 1 (1/0); '?' for anything else (a dot part way from
the middle, a colour that is not JOY's). Half way, as DVX's own joystick
tests judge an axis: an uncalibrated DVX axis ranges itself as it moves,
so an emulated stick's full push reads a little short of the end.

    joyScreen.py PNG [--scale S] [--st] [--at X,Y]

--st: the screen shows the plain Atari ST's 3-bit palette, which drops each
channel's low bit (Hatari shows 3-bit n as the STe's 4-bit 2n). --at: the
stage starts there on the screen (at the scale --scale gives, else 1), so
the patches need not be seen: where something covers one (the PSP's
keyboard), its stick's connected character is '?'.

Prints "A=<state> B=<state> scale=<S> dots=<A's x,y offset> <B's>"; exit 1
when the stage is not found or a stick cannot be read.
"""

import sys

from PIL import Image

# examples/joy/joy.c: its layout (stage pixels) and colours ($0RGB).
STICK_LEFT = (24, 216)
STICK_TOP = 16
FRAME_SIZE = 80
FRAME_INSET = 2
DOT_SIZE = 8
DOT_TRAVEL = (FRAME_SIZE - DOT_SIZE) // 2 - FRAME_INSET
BUTTON_W = 24
BUTTON_H = 16
BUTTON_GAP = 4
ROW_GAP = 6
BUTTON_TOP = STICK_TOP + FRAME_SIZE + ROW_GAP
STATUS_W = 80
STATUS_H = 4
STATUS_TOP = BUTTON_TOP + BUTTON_H + ROW_GAP
COLOURS = {
    "background": 0x000, "frame": 0x444, "dotOff": 0x666, "dotOn": 0x0F0,
    "buttonOff": 0x333, "buttonOn": 0xFF0, "disconnected": 0xF00,
}
# The patch's green is the dot's (COLOR_CONNECTED = COLOR_DOT_ON).
CONNECTED = "dotOn"
# A 15-bit desktop is off by at most 8 a channel, the 6-bit VGA DAC by at
# most 3; JOY's colours sit much further apart. (All are exact in DVX's
# 256-colour desktop palette.)
MAX_DIST = 3 * 9 ** 2
SCALES = (1, 2, 3)
UNKNOWN = "?"


def rgb(v):
    return (((v >> 8) & 0xF) * 17, ((v >> 4) & 0xF) * 17, (v & 0xF) * 17)


# A plain ST's 3-bit palette keeps a channel's top three bits.
ST_CHANNEL_MASK = 0xEEE


def palette(mask):
    return {name: rgb(v & mask) for name, v in COLOURS.items()}


class Screen:
    def __init__(self, img, colours):
        self.w, self.h = img.size
        self.pixels = list(img.getdata())
        self.colours = colours
        self.names = {}

    def name(self, x, y):
        """JOY's colour name at screen (x, y), or None."""
        if not (0 <= x < self.w and 0 <= y < self.h):
            return None
        c = self.pixels[y * self.w + x]
        if c not in self.names:
            best, bestDist = None, MAX_DIST + 1
            for n, p in self.colours.items():
                d = sum((a - b) ** 2 for a, b in zip(c, p))
                if d < bestDist:
                    best, bestDist = n, d
            self.names[c] = best
        return self.names[c]


def findStage(scr, scale):
    """The screen position of stage (0, 0), found by the two status patches:
    runs of STATUS_W stage pixels in either patch colour on one row, as far
    apart as the two frames."""
    run = STATUS_W * scale
    gap = (STICK_LEFT[1] - STICK_LEFT[0]) * scale
    patch = (CONNECTED, "disconnected")
    for y in range(scr.h - STATUS_H * scale + 1):
        x = 0
        while x < scr.w:
            if scr.name(x, y) not in patch:
                x += 1
                continue
            start = x
            while x < scr.w and scr.name(x, y) in patch:
                x += 1
            if x - start != run:
                continue
            other = start + gap
            if all(scr.name(other + i, y) in patch for i in range(run)) and scr.name(other + run, y) not in patch:
                ox = start - STICK_LEFT[0] * scale
                oy = y - STATUS_TOP * scale
                if ox >= 0 and oy >= 0:
                    return ox, oy
    return None


def stagePixel(scr, origin, scale, x, y):
    """JOY's colour name of stage pixel (x, y): its block's middle."""
    return scr.name(origin[0] + x * scale + scale // 2, origin[1] + y * scale + scale // 2)


def axisChar(offset):
    if offset == 0:
        return "0"
    if offset <= -DOT_TRAVEL // 2:
        return "-"
    if offset >= DOT_TRAVEL // 2:
        return "+"
    return UNKNOWN


def stick(scr, origin, scale, left):
    """The five-character state of the stick whose frame starts at left, and
    its dot's offset from the middle ("x,y", "?" when there is no dot)."""
    patch = stagePixel(scr, origin, scale, left + STATUS_W // 2, STATUS_TOP + STATUS_H // 2)
    connected = {CONNECTED: "1", "disconnected": "0"}.get(patch, UNKNOWN)

    dots = [(x, y, stagePixel(scr, origin, scale, x, y))
            for y in range(STICK_TOP, STICK_TOP + FRAME_SIZE)
            for x in range(left, left + FRAME_SIZE)]
    dots = [(x, y, n) for x, y, n in dots if n in ("dotOn", "dotOff")]
    ax = ay = UNKNOWN
    offset = UNKNOWN
    if dots:
        x0 = min(d[0] for d in dots)
        y0 = min(d[1] for d in dots)
        whole = len(dots) == DOT_SIZE * DOT_SIZE and max(d[0] for d in dots) == x0 + DOT_SIZE - 1 and max(d[1] for d in dots) == y0 + DOT_SIZE - 1
        colours = set(d[2] for d in dots)
        dx = x0 - (left + FRAME_SIZE // 2 - DOT_SIZE // 2)
        dy = y0 - (STICK_TOP + FRAME_SIZE // 2 - DOT_SIZE // 2)
        offset = "%+d,%+d" % (dx, dy)
        if whole and len(colours) == 1:
            ax = axisChar(dx)
            ay = axisChar(dy)
            # Green away from the middle, grey in it.
            atRest = ax == "0" and ay == "0"
            if colours != ({"dotOff"} if atRest else {"dotOn"}):
                ax = ay = UNKNOWN

    centre = left + FRAME_SIZE // 2
    buttons = ""
    for bx in (centre - BUTTON_W - BUTTON_GAP // 2, centre + BUTTON_GAP // 2):
        n = stagePixel(scr, origin, scale, bx + BUTTON_W // 2, BUTTON_TOP + BUTTON_H // 2)
        buttons += {"buttonOn": "1", "buttonOff": "0"}.get(n, UNKNOWN)
    return connected + ax + ay + buttons, offset


def main(argv):
    if len(argv) < 2:
        sys.stderr.write(__doc__)
        return 2
    scales = SCALES
    mask = 0xFFF
    at = None
    args = argv[2:]
    while args:
        if args[0] == "--scale" and len(args) > 1:
            scales = (int(args[1]),)
            args = args[2:]
        elif args[0] == "--st":
            mask = ST_CHANNEL_MASK
            args = args[1:]
        elif args[0] == "--at" and len(args) > 1:
            at = tuple(int(v) for v in args[1].split(","))
            args = args[2:]
        else:
            sys.stderr.write(__doc__)
            return 2

    scr = Screen(Image.open(argv[1]).convert("RGB"), palette(mask))
    origin = None
    for scale in scales:
        origin = at if at else findStage(scr, scale)
        if origin:
            break
    if not origin:
        print("joyScreen: FAIL - JOY's two status patches are not on the screen")
        return 1
    read = [stick(scr, origin, scale, left) for left in STICK_LEFT]
    print("A=%s B=%s scale=%d dots=%s %s" % (read[0][0], read[1][0], scale, read[0][1], read[1][1]))
    return 1 if any(UNKNOWN in state for state, offset in read) else 0


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