#!/usr/bin/env python3
"""Drive the qemu monitor: sendkey a string, and screendump to PNG.

Usage:
  vm.py type "some text"     type printable ASCII into the guest
  vm.py key ret tab spc      send named keys
  vm.py shot name            screendump to name.png
"""
import socket, sys, time, zlib, struct, os, re

SCRATCH = os.environ.get(
    "VM_STATE", os.path.join(os.path.dirname(os.path.abspath(__file__)), "state")
)
SOCK = os.path.join(SCRATCH, "monitor.sock")

NAMED = {
    " ": "spc", "-": "minus", "=": "equal", "[": "bracket_left",
    "]": "bracket_right", ";": "semicolon", "'": "apostrophe",
    "`": "grave_accent", "\\": "backslash", ",": "comma", ".": "dot",
    "/": "slash",
}
SHIFTED = {
    "!": "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 Monitor:
    def __init__(self):
        self.s = socket.socket(socket.AF_UNIX)
        self.s.connect(SOCK)
        time.sleep(0.2)
        self.drain()

    def drain(self):
        self.s.setblocking(False)
        out = b""
        try:
            while True:
                chunk = self.s.recv(65536)
                if not chunk:
                    break
                out += chunk
        except BlockingIOError:
            pass
        self.s.setblocking(True)
        return out.decode(errors="replace")

    def cmd(self, line, wait=0.05):
        self.s.sendall((line + "\n").encode())
        time.sleep(wait)
        return self.drain()

    def sendkey(self, k):
        self.cmd("sendkey " + k, wait=0.04)

    def type(self, text):
        for c in text:
            if c.isalnum():
                self.sendkey(c if not c.isupper() else "shift-" + c.lower())
            elif c in NAMED:
                self.sendkey(NAMED[c])
            elif c in SHIFTED:
                self.sendkey("shift-" + SHIFTED[c])
            elif c == "\n":
                self.sendkey("ret")
            else:
                raise SystemExit("no key mapping for %r" % c)


def ppm_to_png(ppm_path, png_path):
    with open(ppm_path, "rb") as f:
        data = f.read()
    # P6 header: magic, width height, maxval, each possibly comment-separated
    fields, idx = [], 2
    while len(fields) < 3:
        while data[idx:idx + 1].isspace():
            idx += 1
        if data[idx:idx + 1] == b"#":
            while data[idx:idx + 1] != b"\n":
                idx += 1
            continue
        start = idx
        while not data[idx:idx + 1].isspace():
            idx += 1
        fields.append(int(data[start:idx]))
    idx += 1
    w, h, _maxval = fields
    pixels = data[idx:]
    raw = b"".join(b"\x00" + pixels[y * w * 3:(y + 1) * w * 3] for y in range(h))

    def chunk(tag, payload):
        return (struct.pack(">I", len(payload)) + tag + payload
                + struct.pack(">I", zlib.crc32(tag + payload) & 0xFFFFFFFF))

    png = (b"\x89PNG\r\n\x1a\n"
           + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
           + chunk(b"IDAT", zlib.compress(raw, 6))
           + chunk(b"IEND", b""))
    with open(png_path, "wb") as f:
        f.write(png)
    return w, h


def main():
    if len(sys.argv) < 2:
        raise SystemExit(__doc__)
    verb = sys.argv[1]
    if verb == "shot":
        name = sys.argv[2]
        ppm = os.path.join(SCRATCH, name + ".ppm")
        png = os.path.join(SCRATCH, name + ".png")
        m = Monitor()
        m.cmd("screendump " + ppm, wait=1.0)
        for _ in range(20):
            if os.path.exists(ppm) and os.path.getsize(ppm) > 1000:
                break
            time.sleep(0.3)
        print(ppm_to_png(ppm, png), png)
        return
    m = Monitor()
    if verb == "type":
        m.type(sys.argv[2])
    elif verb == "key":
        for k in sys.argv[2:]:
            m.sendkey(k)
    elif verb == "hold":
        # Keep the GRUB countdown from expiring: one connection, many keys.
        end = time.time() + float(sys.argv[2])
        while time.time() < end:
            m.sendkey("down")
            m.sendkey("up")
            time.sleep(0.2)
    elif verb == "cmd":
        print(m.cmd(" ".join(sys.argv[2:]), wait=0.5))
    else:
        raise SystemExit(__doc__)


if __name__ == "__main__":
    main()