#!/usr/bin/env python3
"""Minimal QMP driver for the Nomarchy test VMs.

Drives a QEMU instance through its QMP socket: inject keystrokes, type
text, take screendumps, power off. Used by tools/test-install.sh; also
handy interactively while debugging a VM that has no SSH.

Usage:
  qmp.py key <combo>            # e.g. meta_l-ret, ctrl-alt-f2, ret
  qmp.py type <text...>         # types literal text (US layout)
  qmp.py screenshot <out.png>   # NOTE: with virtio-vga-gl this fails
                                # ("no surface") once the guest uses GL —
                                # use vncshot.py instead for desktops
  qmp.py quit                   # hard power-off

The socket path comes from $NOMARCHY_QMP_SOCK
(default /tmp/nomarchy-vm/qmp.sock).

Keystroke pacing: the guest drops keys when busy (especially during
boot). $NOMARCHY_TYPE_DELAY (seconds between keys, default 0.08) trades
speed for reliability; verify long typed lines with a screenshot before
pressing enter.
"""
import json
import os
import socket
import sys
import time

SOCK = os.environ.get("NOMARCHY_QMP_SOCK", "/tmp/nomarchy-vm/qmp.sock")
TYPE_DELAY = float(os.environ.get("NOMARCHY_TYPE_DELAY", "0.08"))

CHAR_KEYS = {
    " ": "spc", "-": "minus", "=": "equal", "[": "bracket_left",
    "]": "bracket_right", ";": "semicolon", "'": "apostrophe",
    "`": "grave_accent", "\\": "backslash", ",": "comma", ".": "dot",
    "/": "slash", "\n": "ret", "\t": "tab",
}
SHIFT_CHARS = {
    "!": "1", "@": "2", "#": "3", "$": "4", "%": "5", "^": "6",
    "&": "7", "*": "8", "(": "9", ")": "0", "_": "minus", "+": "equal",
    "{": "bracket_left", "}": "bracket_right", ":": "semicolon",
    '"': "apostrophe", "~": "grave_accent", "|": "backslash",
    "<": "comma", ">": "dot", "?": "slash",
}


class QMP:
    def __init__(self):
        self.s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.s.connect(SOCK)
        self.f = self.s.makefile("rw")
        self._read()  # greeting
        self.cmd("qmp_capabilities")

    def _read(self):
        while True:
            msg = json.loads(self.f.readline())
            if "event" in msg:
                continue
            return msg

    def cmd(self, name, **args):
        self.f.write(json.dumps({"execute": name, "arguments": args}) + "\n")
        self.f.flush()
        resp = self._read()
        if "error" in resp:
            raise RuntimeError(f"{name}: {resp['error']}")
        return resp

    def send_keys(self, names, hold=40):
        keys = [{"type": "qcode", "data": n} for n in names]
        self.cmd("send-key", keys=keys, **{"hold-time": hold})

    def type_text(self, text):
        for ch in text:
            if ch.isalpha() and ch.lower() != ch:
                self.send_keys(["shift", ch.lower()])
            elif ch in SHIFT_CHARS:
                self.send_keys(["shift", SHIFT_CHARS[ch]])
            elif ch in CHAR_KEYS:
                self.send_keys([CHAR_KEYS[ch]])
            else:
                self.send_keys([ch])
            time.sleep(TYPE_DELAY)


def main():
    q = QMP()
    cmd = sys.argv[1]
    if cmd == "screenshot":
        q.cmd("screendump", filename=sys.argv[2], format="png")
    elif cmd == "key":
        q.send_keys(sys.argv[2].split("-"))
    elif cmd == "type":
        q.type_text(" ".join(sys.argv[2:]))
    elif cmd == "quit":
        q.f.write(json.dumps({"execute": "quit"}) + "\n")
        q.f.flush()
    else:
        sys.exit(f"unknown command {cmd}")


if __name__ == "__main__":
    main()