#!/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 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()
        h = hashlib.sha256(data).hexdigest()
        if h != self.last_hash:
            self.last_hash = h
            self.stable_since = time.time()
        self.img = Image.open(io.BytesIO(data)) if Image else None
        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 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 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
    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')
        if now - start < args.min_boot or screen.stable_for() < args.settle:
            continue
        if screen.is_blank():
            last_progress = now
            continue

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

        # 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

        # once the bootstrap command is typed, only the picker (above) matters
        if typed_at is not 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')
            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):
    screen = Screen(qmp, args.debug_dir, log)
    start = time.time()
    last_periodic = 0
    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)
        if time.time() - last_periodic > args.periodic:
            last_periodic = time.time()
            try:
                screen.grab()
                screen.save('periodic')
            except Exception as e:  # noqa: BLE001
                log(f'screendump failed ({e})')


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('--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())