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

"""qmp.py - drive a QEMU machine through its QMP socket.

The Be API ports' gates (scripts/lib/beGate.sh) take the screen and press keys
this way, as a person at the machine would: the keys arrive through the
emulated PS/2 keyboard.

    qmp.py SOCKET screendump FILE.png   the screen, as QEMU's display has it
    qmp.py SOCKET type TEXT             type TEXT (US layout, shift as needed)
    qmp.py SOCKET keys KEY...           press the KEYs, one after another
    qmp.py SOCKET hold KEY MS           hold one key down for MS milliseconds
    qmp.py SOCKET move DX DY            move the mouse DX, DY (relative, as a PS/2
                                        mouse moves, in steps a PS/2 packet holds;
                                        the system stops it at the screen's edge)
    qmp.py SOCKET quit                  end QEMU

KEY is a QEMU key name (a qcode: esc, ret, spc, a, 1, ...).
"""

import json
import socket
import sys
import time

# How long each typed key is held, and the pause after it: slow enough that
# a program reading its input once a frame sees every key.
TYPE_HOLD_MS = 100
TYPE_GAP_SECONDS = 0.15
MS_PER_SECOND = 1000.0
# A mouse move goes in steps of at most this many counts on each axis, with
# this pause after each: BeOS R5's mouse driver gets a single long move wrong.
MOVE_STEP_MAX = 100
MOVE_GAP_SECONDS = 0.02

# Characters to (qcode, shifted) on a US keyboard.
SHIFTED_DIGITS = ')!@#$%^&*('
PUNCTUATION = {
    ' ': ('spc', False), '.': ('dot', False), ',': ('comma', False),
    '/': ('slash', False), ';': ('semicolon', False), "'": ('apostrophe', False),
    '-': ('minus', False), '=': ('equal', False), '[': ('bracket_left', False),
    ']': ('bracket_right', False), '\\': ('backslash', False), '`': ('grave_accent', False),
    ':': ('semicolon', True), '?': ('slash', True), '>': ('dot', True), '<': ('comma', True),
    '"': ('apostrophe', True), '_': ('minus', True), '+': ('equal', True),
    '{': ('bracket_left', True), '}': ('bracket_right', True), '|': ('backslash', True),
    '~': ('grave_accent', True),
}


class Qmp:
    """One QMP connection: the greeting read and capabilities negotiated."""

    def __init__(self, path):
        self.sock = socket.socket(socket.AF_UNIX)
        self.sock.connect(path)
        self.file = self.sock.makefile('rw')
        self.file.readline()
        self.execute('qmp_capabilities')

    def execute(self, command, arguments=None):
        """Run one command; its return value. Events in between are skipped."""
        request = {'execute': command}
        if arguments is not None:
            request['arguments'] = arguments
        self.file.write(json.dumps(request) + '\n')
        self.file.flush()
        while True:
            reply = json.loads(self.file.readline())
            if 'error' in reply:
                raise RuntimeError('%s: %s' % (command, reply['error'].get('desc', reply['error'])))
            if 'return' in reply:
                return reply['return']

    def keys(self, qcodes, holdMs):
        """Press qcodes together, hold them holdMs, let them go."""
        self.execute('send-key', {
            'keys': [{'type': 'qcode', 'data': code} for code in qcodes],
            'hold-time': holdMs,
        })


def charKey(ch):
    """(qcode, shifted) for one character."""
    if 'a' <= ch <= 'z' or '0' <= ch <= '9':
        return ch, False
    if 'A' <= ch <= 'Z':
        return ch.lower(), True
    if ch in SHIFTED_DIGITS:
        return str(SHIFTED_DIGITS.index(ch)), True
    if ch in PUNCTUATION:
        return PUNCTUATION[ch]
    raise ValueError('no key types %r' % ch)


def move(qmp, dx, dy):
    """Move the mouse dx, dy in MOVE_STEP_MAX steps."""
    steps = max(1, -(-max(abs(dx), abs(dy)) // MOVE_STEP_MAX))
    for step in range(steps):
        stepX = dx * (step + 1) // steps - dx * step // steps
        stepY = dy * (step + 1) // steps - dy * step // steps
        qmp.execute('input-send-event', {'events': [
            {'type': 'rel', 'data': {'axis': 'x', 'value': stepX}},
            {'type': 'rel', 'data': {'axis': 'y', 'value': stepY}},
        ]})
        time.sleep(MOVE_GAP_SECONDS)


def tap(qmp, qcodes):
    """Press qcodes together for a typed key's time, then pause after them."""
    qmp.keys(qcodes, TYPE_HOLD_MS)
    time.sleep(TYPE_HOLD_MS / MS_PER_SECOND + TYPE_GAP_SECONDS)


def main(argv):
    if len(argv) < 3:
        sys.exit(__doc__)
    qmp = Qmp(argv[1])
    command = argv[2]
    if command == 'screendump' and len(argv) == 4:
        qmp.execute('screendump', {'filename': argv[3], 'format': 'png'})
    elif command == 'type' and len(argv) == 4:
        for ch in argv[3]:
            code, shifted = charKey(ch)
            tap(qmp, ['shift', code] if shifted else [code])
    elif command == 'keys' and len(argv) >= 4:
        for code in argv[3:]:
            tap(qmp, [code])
    elif command == 'hold' and len(argv) == 5:
        holdMs = int(argv[4])
        qmp.keys([argv[3]], holdMs)
        time.sleep(holdMs / MS_PER_SECOND)
    elif command == 'move' and len(argv) == 5:
        move(qmp, int(argv[3]), int(argv[4]))
    elif command == 'quit' and len(argv) == 3:
        qmp.execute('quit')
    else:
        sys.exit(__doc__)


if __name__ == '__main__':
    main(sys.argv)