#!/usr/bin/env python3
"""vmix macOS VM driver.

Launches QEMU with a QMP socket and either

  --mode install  drives macOS Recovery to a Terminal with keystrokes (screen
                  settle detection + OCR of the menu bar), types the bootstrap
                  command and waits for the VM to power itself off
  --mode boot     waits for the VM to power itself off (customize steps)

Everything after `--` is the QEMU command line. Screenshots and a log are
written to --debug-dir (default /tmp/vmix-macos/<name>) for troubleshooting.
"""
import argparse
import hashlib
import io
import json
import os
import socket
import subprocess
import sys
import time

try:
    from PIL import Image
except ImportError:  # pragma: no cover
    Image = None
try:
    import pytesseract
except ImportError:  # pragma: no cover
    pytesseract = None


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


class Log:
    def __init__(self, path):
        self.f = open(path, 'a')
        self.t0 = time.time()

    def __call__(self, msg):
        line = f'[{time.time() - self.t0:7.1f}s] {msg}'
        print(f'vmix driver: {line}', flush=True)
        self.f.write(line + '\n')
        self.f.flush()


class QMP:
    def __init__(self, path):
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(path)
        self.f = self.sock.makefile('rwb', buffering=0)
        self._read()
        self.cmd('qmp_capabilities')

    def _read(self):
        while True:
            line = self.f.readline()
            if not line:
                raise EOFError('QMP connection closed')
            msg = json.loads(line)
            if 'event' in msg:
                continue
            return msg

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

    def screendump(self, path):
        self.cmd('screendump', filename=path)

    def system_reset(self):
        self.cmd('system_reset')

    def system_powerdown(self):
        self.cmd('system_powerdown')

    _jig = 0

    def jiggle(self):
        # tiny absolute (usb-tablet) pointer move to keep the display awake: a real
        # HID event, but < 2 screen px so it does not change the settle fingerprint
        self._jig = 16060 if self._jig < 16030 else 16000
        try:
            self.cmd('input-send-event', events=[
                {'type': 'abs', 'data': {'axis': 'x', 'value': self._jig}},
                {'type': 'abs', 'data': {'axis': 'y', 'value': 16000}}])
        except Exception:  # noqa: BLE001
            self.send_key('shift')

    def send_key(self, *keys, hold=80):
        self.cmd('send-key', keys=[{'type': 'qcode', 'data': k} for k in keys], **{'hold-time': hold})
        time.sleep(0.15)

    def type_text(self, text):
        for ch in text:
            if ch.isascii() and ch.isalnum():
                if ch.isupper():
                    self.send_key('shift', ch.lower())
                else:
                    self.send_key(ch)
            elif ch in PLAIN:
                self.send_key(PLAIN[ch])
            elif ch in SHIFTED:
                self.send_key('shift', SHIFTED[ch])
            else:
                raise ValueError(f'cannot type {ch!r}')


class Screen:
    """Screenshot helper: settle detection via hashing, OCR of regions."""

    def __init__(self, qmp, debug_dir, log):
        self.qmp = qmp
        self.debug_dir = debug_dir
        self.log = log
        import tempfile
        fd, self.tmp = tempfile.mkstemp(prefix='.shot-', suffix='.ppm', dir=debug_dir)
        os.close(fd)
        os.chmod(self.tmp, 0o666)
        self.n = 0
        self.last_hash = None
        self.stable_since = time.time()
        self.img = None

    def grab(self):
        self.qmp.screendump(self.tmp)
        with open(self.tmp, 'rb') as f:
            data = f.read()
        self.img = Image.open(io.BytesIO(data)) if Image else None
        # fingerprint from a coarse, quantized grayscale thumbnail so the moving
        # mouse cursor (keepalive jiggle) does not count as a screen change
        if self.img is not None:
            px = self.img.convert('L').resize((48, 36))
            fp = bytes(b & 0xF0 for b in px.getdata())
            h = hashlib.sha256(fp).hexdigest()
        else:
            h = hashlib.sha256(data).hexdigest()
        if h != self.last_hash:
            self.last_hash = h
            self.stable_since = time.time()
        return self.img

    def stable_for(self):
        return time.time() - self.stable_since

    def save(self, tag):
        self.n += 1
        path = os.path.join(self.debug_dir, f'{self.n:03d}-{tag}.png')
        try:
            if self.img is not None:
                self.img.save(path)
            else:
                os.link(self.tmp, path.replace('.png', '.ppm'))
        except Exception as e:  # noqa: BLE001
            self.log(f'could not save screenshot: {e}')
        return path

    def ocr(self, region=None, scale=3, psm=6):
        if self.img is None or pytesseract is None:
            return ''
        img = self.img
        if region:
            img = img.crop(region)
        img = img.convert('L').resize((img.width * scale, img.height * scale), Image.LANCZOS)
        try:
            return pytesseract.image_to_string(img, config=f'--psm {psm}').lower()
        except Exception as e:  # noqa: BLE001
            self.log(f'ocr failed: {e}')
            return ''

    def mean(self):
        if self.img is None:
            return 128
        px = self.img.convert('L').resize((32, 24))
        d = list(px.getdata())
        return sum(d) / len(d)

    def is_blank(self):
        # black/uniform screen (firmware, boot): nothing to act on
        if self.img is None:
            return False
        lo, hi = self.img.convert('L').resize((64, 48)).getextrema()
        return hi - lo < 24

    def menubar_text(self):
        w = self.img.width if self.img else 1024
        return self.ocr((0, 0, w, 40), scale=4, psm=7)


def prepare_debug_dir(path):
    # nix builds run as different nixbld users: keep the shared dirs world-writable
    try:
        for d in (os.path.dirname(path), path):
            os.makedirs(d, exist_ok=True)
            try:
                os.chmod(d, 0o1777 if d != path else 0o777)
            except OSError:
                pass
        probe = os.path.join(path, '.probe')
        open(probe, 'w').close()
        os.unlink(probe)
        # a rebuild reuses this dir but runs as a different nixbld user; drop stale
        # files so screendumps/PNGs are not blocked by another owner's 0644 files
        import glob
        for f in glob.glob(os.path.join(path, '*')) + glob.glob(os.path.join(path, '.current*')):
            try:
                os.unlink(f)
            except OSError:
                pass
        return path
    except OSError:
        import tempfile
        alt = tempfile.mkdtemp(prefix='vmix-macos-')
        print(f'vmix driver: {path} not writable, using {alt}', flush=True)
        return alt


def launch(qemu_args, qmp_sock, log):
    if os.path.exists(qmp_sock):
        os.unlink(qmp_sock)
    args = list(qemu_args) + ['-qmp', f'unix:{qmp_sock},server,nowait']
    log('launching: ' + ' '.join(args))
    proc = subprocess.Popen(args)
    deadline = time.time() + 60
    while not os.path.exists(qmp_sock):
        if proc.poll() is not None:
            return proc, None
        if time.time() > deadline:
            proc.kill()
            raise RuntimeError('QEMU did not create the QMP socket')
        time.sleep(0.2)
    time.sleep(0.5)
    return proc, QMP(qmp_sock)


def open_terminal(qmp, log):
    # Ctrl-F2 focuses the menu bar; typing jumps to the menu whose title starts
    # with that letter (Utilities), Down opens it, "t" jumps to Terminal.
    log('opening Terminal via menu bar (ctrl-f2, u, down, t, ret)')
    qmp.send_key('ctrl', 'f2')
    time.sleep(1.0)
    qmp.send_key('u')
    time.sleep(0.7)
    qmp.send_key('down')
    time.sleep(0.7)
    qmp.send_key('t')
    time.sleep(0.7)
    qmp.send_key('ret')


def disk_idle(args):
    """True if the system disk has had no writes recently (guest not doing I/O)."""
    if not args.progress_file:
        return True
    try:
        return (time.time() - os.path.getmtime(args.progress_file)) > args.disk_idle
    except OSError:
        return True


def run_install(args, proc, qmp, log):
    """Drive the install VM to completion.

    OpenCore shows a boot picker on every (re)boot and does not always auto-boot,
    so on any settled picker we press Return to boot the highlighted macOS entry
    (aux entries are hidden; during the install phases startosinstall blesses the
    right default). That runs on EVERY iteration, because the install reboots
    several times after we hand off to startosinstall. Before we have typed the
    bootstrap command we also drive Recovery: language/welcome -> Return, the
    Recovery window -> open Terminal, Terminal -> type the command.
    """
    RECOVERY_BODY = ('reinstall', 'disk utility', 'restore from', 'recovery assistant',
                     'macos utilities')
    PICKER_BODY = ('base system', 'macos installer', 'rel-1', 'rel-0')  # OpenCore picker
    LANG_BODY = ('language', 'select your', 'main language', 'country or region',
                 'welcome', 'get started', 'choose your')
    screen = Screen(qmp, args.debug_dir, log)
    start = time.time()
    typed_at = None
    terminal_attempts = 0
    last_periodic = 0
    last_progress = start
    blind_done = False
    resets = 0
    blank_since = None
    login_since = None
    recovery_start = None
    last_term_action = 0
    while True:
        rc = proc.poll()
        if rc is not None:
            return rc
        now = time.time()
        if now - start > args.timeout:
            try:
                screen.grab(); screen.save('timeout')
            except Exception:  # noqa: BLE001
                pass
            log('timeout reached, killing QEMU')
            proc.kill()
            return 124
        time.sleep(args.interval)
        try:
            screen.grab()
        except Exception as e:  # noqa: BLE001
            log(f'screendump failed ({e}), assuming QEMU is exiting')
            time.sleep(2)
            continue
        if now - last_periodic > args.periodic:
            last_periodic = now
            screen.save('periodic')
        # keep the recovery display awake until the command is typed (mouse jiggle)
        if typed_at is None and now - last_term_action > 8:
            qmp.jiggle()
        if now - start < args.min_boot or screen.stable_for() < args.settle:
            continue
        if screen.is_blank():
            last_progress = now
            if blank_since is None:
                blank_since = now
            if typed_at is None:
                # recovery display asleep — jiggle the mouse to wake it, wait for UI
                qmp.jiggle()
                continue
            # macOS `shutdown -h now` halts the guest to a black screen without an
            # ACPI power-off, so QEMU never exits. Once we have handed off (command
            # typed), a long pure-black screen means the agent finished and halted.
            elif typed_at is not None and now - blank_since > args.halt_timeout and disk_idle(args):
                screen.save('halt')
                log(f'guest halted (black {now - blank_since:.0f}s, disk idle); killing QEMU, readback will validate')
                proc.kill()
                try:
                    proc.wait(timeout=10)
                except Exception:  # noqa: BLE001
                    pass
                return 0
            continue
        blank_since = None

        top = screen.menubar_text()
        body = screen.ocr()
        log(f'settled: menubar={top.strip()!r} body~={" ".join(body.split())[:80]!r}')

        # Boot-hang watchdog: a dark screen (Apple logo / black) frozen for a long
        # time with no menu bar is a stuck (re)boot — kick it with a system reset.
        # Never fires on the bright, static Terminal of the prepare phase.
        if 'terminal' not in top and 'utilities' not in top and screen.mean() < 40 \
                and screen.stable_for() > args.stall_reset and disk_idle(args) and resets < args.max_resets:
            resets += 1
            screen.save('stall-reset')
            log(f'boot hung ({screen.stable_for():.0f}s frozen, dark, disk idle), system_reset #{resets}')
            try:
                qmp.system_reset()
            except Exception as e:  # noqa: BLE001
                log(f'system_reset failed: {e}')
            screen.stable_since = time.time()
            last_progress = now
            continue

        # OpenCore boot picker — always handle it (the install reboots many times)
        if 'terminal' not in top and 'utilities' not in top and any(k in body for k in PICKER_BODY):
            screen.save('picker')
            log('OpenCore boot picker, pressing Return to boot the default macOS entry')
            qmp.send_key('ret')
            last_progress = now
            screen.stable_since = time.time()
            continue

        # After the install, the loginwindow/desktop is a BRIGHT gray screen, unlike
        # the dark install/boot screens (Apple logo). The vmix agent powers the VM
        # off if it runs (cron/daemon); if BTM blocks it, we power down here so the
        # build still completes with a bootable, installed image. Brightness is a
        # far more reliable signal than OCR of the faint "password" text.
        bright = screen.mean() > 80
        loginish = (typed_at is not None and bright and 'terminal' not in top
                    and 'utilities' not in top and not any(k in body for k in PICKER_BODY))
        if loginish:
            if login_since is None:
                login_since = now
                log('bright post-install screen (loginwindow/desktop) — OS installed; grace before powerdown')
            elif now - login_since > args.login_grace:
                screen.save('loginwindow')
                log(f'loginwindow persisted {now - login_since:.0f}s, powering down (install complete)')
                try:
                    qmp.system_powerdown()
                except Exception as e:  # noqa: BLE001
                    log(f'powerdown failed: {e}')
                for _ in range(90):
                    if proc.poll() is not None:
                        return 0
                    time.sleep(1)
                proc.kill()
                return 0
            continue
        else:
            login_since = None

        # once the bootstrap command is typed, only the picker (above) and an
        # unexpected return to Recovery matter (post-prepare reboot landed on the
        # recovery instead of the installer — restart the install then).
        if typed_at is not None:
            if ('utilities' in top or 'recovery' in top):
                if recovery_start is None:
                    recovery_start = now
                if now - typed_at > 120 and now - recovery_start > 45:
                    log('unexpectedly back at Recovery after install started — restarting install')
                    typed_at = None
                    terminal_attempts = 0
                    recovery_start = None
                    # fall through to the recovery/terminal handling below
                else:
                    continue
            else:
                recovery_start = None
                continue

        if 'terminal' in top:
            screen.save('terminal')
            log(f'typing bootstrap command: {args.command!r}')
            qmp.type_text(args.command + '\n')
            typed_at = time.time()
            continue

        acted = False
        if 'utilities' in top or 'recovery' in top or any(k in body for k in RECOVERY_BODY):
            screen.save('recovery')
            terminal_attempts += 1
            log(f'recovery window (attempt {terminal_attempts}), opening Terminal')
            last_term_action = now
            open_terminal(qmp, log)
            if terminal_attempts >= 3:
                time.sleep(8)
                log('typing bootstrap command (Terminal assumed open)')
                qmp.type_text(args.command + '\n')
                typed_at = time.time()
                continue
            acted = True
        elif any(k in body for k in LANG_BODY):
            screen.save('language')
            log('language/welcome screen, pressing Return')
            qmp.send_key('ret')
            acted = True

        if acted:
            last_progress = now
            screen.stable_since = time.time()
        elif now - last_progress > args.settle * args.max_actions and not blind_done:
            blind_done = True
            screen.save('blind')
            log('nothing recognised for a long time, blind sequence')
            qmp.send_key('ret')
            time.sleep(20)
            open_terminal(qmp, log)
            time.sleep(10)
            qmp.type_text(args.command + '\n')
            typed_at = time.time()


def run_boot(args, proc, qmp, log):
    """Wait for the VM to power itself off (customize/generalize/first-boot),
    handling the OpenCore picker and kicking a hung boot with a system reset."""
    PICKER_BODY = ('base system', 'macos installer', 'macintosh hd', 'rel-1', 'rel-0')
    screen = Screen(qmp, args.debug_dir, log)
    start = time.time()
    last_periodic = 0
    resets = 0
    blank_since = None
    login_since = None
    while True:
        rc = proc.poll()
        if rc is not None:
            return rc
        if time.time() - start > args.timeout:
            try:
                screen.grab(); screen.save('timeout')
            except Exception:  # noqa: BLE001
                pass
            log('timeout reached, killing QEMU')
            proc.kill()
            return 124
        time.sleep(args.interval)
        try:
            screen.grab()
        except Exception as e:  # noqa: BLE001
            log(f'screendump failed ({e})')
            time.sleep(2)
            continue
        now = time.time()
        if now - last_periodic > args.periodic:
            last_periodic = now
            screen.save('periodic')
        if now - start < args.min_boot or screen.stable_for() < args.settle:
            continue
        if screen.is_blank():
            if blank_since is None:
                blank_since = now
            elif now - blank_since > args.halt_timeout and disk_idle(args):
                screen.save('halt')
                log(f'guest halted (black {now - blank_since:.0f}s, disk idle); killing QEMU')
                proc.kill()
                try:
                    proc.wait(timeout=10)
                except Exception:  # noqa: BLE001
                    pass
                return 0
            continue
        blank_since = None
        top = screen.menubar_text()
        body = screen.ocr()
        if 'terminal' not in top and 'utilities' not in top and any(k in body for k in PICKER_BODY):
            screen.save('picker')
            log('OpenCore boot picker, pressing Return')
            qmp.send_key('ret')
            screen.stable_since = time.time()
            login_since = None
            continue
        # bright post-boot screen (loginwindow/desktop) => booted; power down if the
        # agent did not (so customize/generalize completes even if BTM blocks it)
        if screen.mean() > 80 and 'terminal' not in top and 'utilities' not in top:
            if login_since is None:
                login_since = now
                log('bright screen (loginwindow/desktop) after boot; grace before powerdown')
            elif now - login_since > args.login_grace:
                screen.save('loginwindow')
                log(f'loginwindow persisted {now - login_since:.0f}s, powering down')
                try:
                    qmp.system_powerdown()
                except Exception as e:  # noqa: BLE001
                    log(f'powerdown failed: {e}')
                for _ in range(90):
                    if proc.poll() is not None:
                        return 0
                    time.sleep(1)
                proc.kill()
                return 0
            continue
        else:
            login_since = None
        if 'terminal' not in top and 'utilities' not in top and screen.mean() < 40 \
                and screen.stable_for() > args.stall_reset and disk_idle(args) and resets < args.max_resets:
            resets += 1
            screen.save('stall-reset')
            log(f'boot hung ({screen.stable_for():.0f}s frozen, dark, disk idle), system_reset #{resets}')
            try:
                qmp.system_reset()
            except Exception as e:  # noqa: BLE001
                log(f'system_reset failed: {e}')
            screen.stable_since = time.time()


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--mode', choices=['install', 'boot'], required=True)
    p.add_argument('--name', default='macos')
    p.add_argument('--debug-dir', default=None)
    p.add_argument('--timeout', type=int, default=4 * 3600, help='seconds before QEMU is killed')
    p.add_argument('--interval', type=float, default=5.0, help='seconds between screenshots')
    p.add_argument('--periodic', type=float, default=120.0, help='seconds between saved debug screenshots')
    p.add_argument('--settle', type=float, default=12.0, help='seconds a screen must be unchanged to act on it')
    p.add_argument('--min-boot', type=float, default=45.0, help='seconds before the first action')
    p.add_argument('--max-actions', type=int, default=8)
    p.add_argument('--stall-reset', type=float, default=360.0, help='reset the VM if a non-Terminal screen is frozen this long (boot hang)')
    p.add_argument('--max-resets', type=int, default=6)
    p.add_argument('--halt-timeout', type=float, default=150.0, help='after the bootstrap, a pure-black screen this long means the guest halted (macOS shutdown does not ACPI-power-off QEMU)')
    p.add_argument('--progress-file', default=None, help='a file (the system disk) whose mtime shows guest activity; resets/halt only fire when it is also idle, so a slow-but-working boot is never interrupted')
    p.add_argument('--disk-idle', type=float, default=90.0, help='seconds of no writes to --progress-file that count as idle')
    p.add_argument('--login-grace', type=float, default=240.0, help='seconds to wait at the loginwindow for the agent to power off before the driver powers down itself')
    p.add_argument('--command', default='diskutil mount VMIX;sh /Volumes/VMIX/run.sh')
    p.add_argument('qemu', nargs=argparse.REMAINDER)
    args = p.parse_args()
    qemu_args = args.qemu[1:] if args.qemu and args.qemu[0] == '--' else args.qemu
    if not qemu_args:
        p.error('QEMU command line required after --')

    args.debug_dir = prepare_debug_dir(args.debug_dir or f'/tmp/vmix-macos/{args.name}')
    log = Log(os.path.join(args.debug_dir, 'driver.log'))
    log(f'mode={args.mode} debug-dir={args.debug_dir} ocr={"yes" if pytesseract else "no"}')
    qmp_sock = os.path.join(args.debug_dir, 'qmp.sock')

    proc, qmp = launch(qemu_args, qmp_sock, log)
    if qmp is None and '-display' in qemu_args and 'sdl' in qemu_args:
        # SDL could not open a window: retry headless
        log(f'QEMU exited early ({proc.returncode}) with SDL, retrying headless')
        i = qemu_args.index('-display')
        qemu_args = qemu_args[:i] + ['-display', 'none'] + qemu_args[i + 2:]
        proc, qmp = launch(qemu_args, qmp_sock, log)
    if qmp is None:
        log(f'QEMU exited immediately with {proc.returncode}')
        return proc.returncode or 1

    try:
        if args.mode == 'install':
            rc = run_install(args, proc, qmp, log)
        else:
            rc = run_boot(args, proc, qmp, log)
    finally:
        if proc.poll() is None:
            proc.kill()
    log(f'QEMU exited with {rc}')
    return rc


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