#!/usr/bin/env python3
"""vmix macOS VM driver: runs QEMU and decides when a build boot is finished.

Modes
  pe       the recovery PE runs /Volumes/VMIX/run.sh and powers off. Success is
           QEMU exiting on its own; the caller checks vmix-run.status.
  install  the PE starts the macOS installer, which reboots through its phases
           into the installed system. Finished when that system reaches the
           (bright) loginwindow / Setup Assistant — the driver powers it down —
           or halts on its own.
  boot     boot an installed image and wait for it to halt or reach the loginwindow.
  qga      boot an installed image with Apple's QEMU guest agent attached
           (--qga-sock), wait for it, run --qga-command as root through it
           (guest-exec), then shut the guest down. Used for online templates.

Observation is passive: the serial console (boot-args serial=3: the PE's
"VMIX-*" markers, kernel boots, panics) and screenshots over QMP (mean
brightness + a coarse change fingerprint). No OCR, no keystrokes. A boot hang
(dark, frozen screen, disk and serial idle) is retried with a system_reset.
Screenshots, the driver log and the serial log are kept in --debug-dir.
"""
import argparse
import hashlib
import io
import json
import os
import shutil
import socket
import subprocess
import sys
import tempfile
import time

try:
    from PIL import Image
except ImportError:  # screenshots then only serve as debug files
    Image = None

PANIC_MARKS = ('panic(cpu', 'Kernel Extensions in backtrace', 'Debugger called: <panic>', 'Nested panic detected', 'panic string:')
REBOOT_MARK = 'MACH Reboot'


class Log:
    def __init__(self, path):
        self.f = open(path, 'a') if path else None
        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)
        if self.f:
            self.f.write(line + '\n')
            self.f.flush()


class QMP:
    def __init__(self, path):
        self.path = path
        self.s = None
        self.buf = b''

    def connect(self, timeout=90):
        t0 = time.time()
        while True:
            try:
                s = socket.socket(socket.AF_UNIX)
                s.settimeout(60)
                s.connect(self.path)
                self.s = s
                self.buf = b''
                self._read()            # greeting
                self.cmd('qmp_capabilities')
                return
            except (OSError, ValueError):
                if time.time() - t0 > timeout:
                    raise
                time.sleep(1)

    def _read(self):
        while b'\n' not in self.buf:
            d = self.s.recv(65536)
            if not d:
                raise OSError('QMP socket closed')
            self.buf += d
        line, self.buf = self.buf.split(b'\n', 1)
        return json.loads(line)

    def cmd(self, name, **args):
        self.s.sendall(json.dumps({'execute': name, 'arguments': args}).encode() + b'\n')
        while True:
            r = self._read()
            if 'return' in r:
                return r['return']
            if 'error' in r:
                raise RuntimeError(r['error'])

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


class QGA:
    """Minimal QEMU guest agent client over the unix socket QEMU serves."""

    def __init__(self, path):
        self.path = path

    def cmd(self, name, timeout=30, **args):
        s = socket.socket(socket.AF_UNIX)
        s.settimeout(timeout)
        try:
            s.connect(self.path)
            s.sendall((json.dumps({'execute': name, 'arguments': args}) + '\n').encode())
            buf = b''
            while b'\n' not in buf:
                d = s.recv(65536)
                if not d:
                    raise OSError('guest agent closed the connection')
                buf += d
        finally:
            s.close()
        r = json.loads(buf.split(b'\n', 1)[0])
        if 'error' in r:
            raise RuntimeError(r['error'])
        return r.get('return')

    def ping(self):
        try:
            return self.cmd('guest-ping', timeout=5) == {}
        except Exception:  # noqa: BLE001
            return False

    def exec_start(self, command):
        return self.cmd('guest-exec', path='/bin/bash', arg=['-c', command], **{'capture-output': True})['pid']

    def exec_status(self, pid):
        return self.cmd('guest-exec-status', pid=pid)


class Screen:
    """Screenshots over QMP with a coarse change fingerprint (cursor-insensitive)."""

    def __init__(self, qmp, debug_dir, log):
        self.qmp = qmp
        self.dir = debug_dir
        self.log = log
        self.tmp = os.path.join(debug_dir, f'.grab-{os.getpid()}.ppm')
        self.img = None
        self.last_fp = None
        self.stable_since = time.time()
        self.n = 0

    def grab(self):
        self.qmp.screendump(self.tmp)
        with open(self.tmp, 'rb') as f:
            data = f.read()
        if Image is None:
            fp = hashlib.sha256(data).hexdigest()
        else:
            self.img = Image.open(io.BytesIO(data))
            small = self.img.convert('L').resize((48, 36))
            fp = hashlib.sha256(bytes(b & 0xF0 for b in small.tobytes())).hexdigest()
        if fp != self.last_fp:
            self.last_fp = fp
            self.stable_since = time.time()
        return self.img

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

    def mean(self):
        if self.img is None:
            return 0
        g = self.img.convert('L').resize((64, 48))
        px = g.tobytes()
        return sum(px) / len(px)

    def is_blank(self):
        return self.img is not None and self.mean() < 3

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


class Serial:
    """Tail the serial console file QEMU writes (-serial file:...)."""

    def __init__(self, path):
        self.path = path
        self.pos = 0
        self.last_activity = time.time()
        self.boots = 0
        self.reboot_at = None

    def poll(self):
        if not self.path or not os.path.exists(self.path):
            return []
        with open(self.path, 'rb') as f:
            f.seek(self.pos)
            data = f.read()
            self.pos = f.tell()
        if not data:
            return []
        self.last_activity = time.time()
        lines = data.decode('utf-8', 'replace').replace('\r', '').split('\n')
        for l in lines:
            if l.startswith('Darwin Kernel Version'):
                self.boots += 1
                self.reboot_at = None
            elif REBOOT_MARK in l:
                self.reboot_at = time.time()
        return lines

    def idle_for(self):
        return time.time() - self.last_activity


def prepare_debug_dir(path):
    """Create the debug dir; builds run as different nixbld users, so the parent
    is made world-writable and a temp dir is used if the path is not writable."""
    parent = os.path.dirname(path)
    try:
        if not os.path.isdir(parent):
            os.makedirs(parent, exist_ok=True)
            os.chmod(parent, 0o777)
        os.makedirs(path, exist_ok=True)
        os.chmod(path, 0o777)
    except OSError:
        path = tempfile.mkdtemp(prefix=os.path.basename(path) + '-', dir='/tmp')
        os.chmod(path, 0o777)
    for f in os.listdir(path):
        if f.endswith(('.png', '.ppm', '.log')) or f.startswith('.grab-') or f == 'qmp.sock':
            try:
                os.remove(os.path.join(path, f))
            except OSError:
                pass
    return path


def launch(qemu_args, qmp_sock, log):
    if os.path.exists(qmp_sock):
        os.remove(qmp_sock)
    cmd = list(qemu_args) + ['-qmp', f'unix:{qmp_sock},server,nowait']
    log('launching: ' + ' '.join(cmd))
    return subprocess.Popen(cmd)


def disk_idle(args):
    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_qga(args, proc, qmp, screen, serial, log):
    """Online template: wait for the guest agent, run the command, shut down."""
    import base64
    qga = QGA(args.qga_sock)
    start = time.time()
    last_periodic = 0
    while not qga.ping():
        rc = proc.poll()
        if rc is not None:
            log(f'QEMU exited with {rc} before the guest agent came up')
            return rc or 3
        now = time.time()
        if now - start > args.start_timeout:
            try:
                screen.grab()
                screen.save('no-agent')
            except Exception:  # noqa: BLE001
                pass
            log(f'guest agent not reachable within {args.start_timeout:.0f}s')
            proc.kill()
            return 3
        for line in serial.poll():
            if any(m in line for m in PANIC_MARKS):
                log('serial: ' + line.strip()[:200])
        if now - last_periodic > args.periodic:
            last_periodic = now
            try:
                screen.grab()
                screen.save('periodic')
            except Exception:  # noqa: BLE001
                pass
        time.sleep(3)
    log(f'guest agent up after {time.time() - start:.0f}s: {qga.cmd("guest-info").get("version")}')
    time.sleep(args.qga_settle)      # let the login session / volumes settle
    try:
        pid = qga.exec_start(args.qga_command)
    except Exception as e:  # noqa: BLE001
        log(f'guest-exec failed: {e}')
        proc.kill()
        return 3
    log(f'running command as root (pid {pid}): {args.qga_command[:160]}')
    t0 = time.time()
    while True:
        rc = proc.poll()
        if rc is not None:
            log(f'QEMU exited with {rc} while the command was running')
            return rc or 3
        if time.time() - start > args.timeout:
            log('timeout reached, killing QEMU')
            proc.kill()
            return 124
        try:
            st = qga.exec_status(pid)
        except Exception as e:  # noqa: BLE001
            log(f'guest-exec-status failed: {e}')
            time.sleep(5)
            continue
        if st.get('exited'):
            break
        now = time.time()
        if now - last_periodic > args.periodic:
            last_periodic = now
            try:
                screen.grab()
                screen.save('periodic')
            except Exception:  # noqa: BLE001
                pass
        time.sleep(3)
    out = base64.b64decode(st.get('out-data', '')).decode('utf-8', 'replace')
    err = base64.b64decode(st.get('err-data', '')).decode('utf-8', 'replace')
    code = st.get('exitcode', st.get('signal'))
    log(f'command finished in {time.time() - t0:.0f}s, exit {code}')
    for line in (out + err).splitlines()[-200:]:
        log('guest: ' + line[:220])
    try:
        screen.grab()
        screen.save('after-command')
    except Exception:  # noqa: BLE001
        pass
    log('shutting the guest down')
    try:
        qga.exec_start('sync; /sbin/shutdown -h now')
    except Exception as e:  # noqa: BLE001
        log(f'guest shutdown failed ({e}), ACPI powerdown')
        try:
            qmp.system_powerdown()
        except Exception:  # noqa: BLE001
            pass
    for _ in range(180):
        if proc.poll() is not None:
            log('QEMU exited')
            return 0 if code == 0 else 4
        time.sleep(1)
    log('guest did not power off, killing QEMU')
    proc.kill()
    return 0 if code == 0 else 4


def drive(args, proc, qmp, screen, serial, log):
    start = time.time()
    last_periodic = 0
    resets = 0
    panics = 0
    started = args.mode == 'boot'      # pe/install: wait for the PE marker first
    blank_since = None
    login_since = None
    while True:
        rc = proc.poll()
        if rc is not None:
            log(f'QEMU exited with {rc}')
            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)

        panic = False
        for line in serial.poll():
            if 'VMIX' in line:
                log('serial: ' + line.strip()[:220])
                if 'VMIX-PE: running run.sh' in line and not started:
                    started = True
                    log('PE started run.sh')
                if 'VMIX-PE: no VMIX volume' in line and args.mode != 'boot':
                    log('PE did not find the VMIX volume')
                    proc.kill()
                    return 3
            if line.startswith('Darwin Kernel Version'):
                log(f'guest kernel boot #{serial.boots}')
            if any(m in line for m in PANIC_MARKS):
                panic = True
                log('serial: ' + line.strip()[:220])
        if panic:
            panics += 1
            try:
                screen.grab()
                screen.save('panic')
            except Exception:  # noqa: BLE001
                pass
            if panics > args.max_resets:
                log(f'kernel panic #{panics}, giving up')
                proc.kill()
                return 3
            # XNU reboots by itself after a panic; only reset if no kernel comes back.
            # (In pe mode the PE then runs run.sh again — templates are idempotent.)
            log(f'kernel panic #{panics}, waiting for the guest to reboot')
            serial.reboot_at = now
            continue

        # The guest asked for a reboot but no kernel came back: macOS' restart
        # path panics in QEMU (AppleSMC watchdog keys, see README); reset now
        # instead of waiting for the frozen-screen watchdog.
        if serial.reboot_at and now - serial.reboot_at > args.reboot_timeout:
            resets += 1
            try:
                screen.grab()
                screen.save('reboot-dead')
            except Exception:  # noqa: BLE001
                pass
            log(f'guest requested a reboot {now - serial.reboot_at:.0f}s ago and died, system_reset #{resets}')
            serial.reboot_at = None
            if resets > args.max_resets:
                proc.kill()
                return 3
            qmp.system_reset()
            screen.stable_since = time.time()
            continue

        if not started and now - start > args.start_timeout:
            try:
                screen.grab()
                screen.save('no-start')
            except Exception:  # noqa: BLE001
                pass
            log(f'PE did not start run.sh within {args.start_timeout:.0f}s')
            proc.kill()
            return 3

        try:
            screen.grab()
        except Exception as e:  # noqa: BLE001
            log(f'screendump failed ({e})')
            time.sleep(2)
            continue
        if now - last_periodic > args.periodic:
            last_periodic = now
            screen.save('periodic')

        if args.mode == 'pe':
            continue                    # the PE powers off by itself; nothing to decide
        # install: the PE phase (kernel boot #1) is protected by the guest's own
        # retries; the checks below apply once the installer has rebooted.
        in_os = args.mode == 'boot' or serial.boots >= 2
        if not in_os or screen.stable_for() < args.settle:
            continue
        idle = disk_idle(args) and serial.idle_for() > args.disk_idle

        if screen.is_blank():
            blank_since = blank_since or now
            # macOS `shutdown -h` halts to a black screen without an ACPI power-off
            if now - blank_since > args.halt_timeout and idle:
                screen.save('halt')
                log(f'guest halted (black {now - blank_since:.0f}s, idle); killing QEMU')
                proc.kill()
                return 0
            continue
        blank_since = None

        if screen.mean() > args.bright:
            # loginwindow / Setup Assistant: the OS is installed and booted
            if login_since is None:
                login_since = now
                log(f'bright screen (mean {screen.mean():.0f}): loginwindow/desktop, grace {args.login_grace:.0f}s')
            elif now - login_since > args.login_grace:
                screen.save('loginwindow')
                log('powering down (boot complete)')
                try:
                    qmp.system_powerdown()
                except Exception as e:  # noqa: BLE001
                    log(f'powerdown failed: {e}')
                for _ in range(120):
                    if proc.poll() is not None:
                        log('QEMU exited after powerdown')
                        return 0
                    time.sleep(1)
                # macOS ignores the power button at the Setup Assistant; the
                # volumes are journaled (APFS) and the PE mounts them cleanly next
                log('guest ignores the ACPI power button here (Setup Assistant); stopping QEMU')
                proc.kill()
                return 0
            continue
        login_since = None

        # dark, frozen, nothing happening: a boot hang (seen at the Apple logo)
        if screen.stable_for() > args.stall_reset and idle and resets < args.max_resets:
            resets += 1
            screen.save('stall-reset')
            log(f'boot hung ({screen.stable_for():.0f}s frozen, dark, 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()
    p.add_argument('--mode', choices=['pe', 'install', 'boot', 'qga'], required=True)
    p.add_argument('--qga-sock', default=None, help='guest agent unix socket (mode qga)')
    p.add_argument('--qga-command', default=None, help='bash command to run as root through the guest agent (mode qga)')
    p.add_argument('--qga-settle', type=float, default=20.0, help='seconds to wait after the agent answers before running the command')
    p.add_argument('--name', default='macos')
    p.add_argument('--debug-dir', default=None)
    p.add_argument('--serial-log', default=None, help='file QEMU writes the serial console to')
    p.add_argument('--timeout', type=int, default=4 * 3600, help='seconds before QEMU is killed')
    p.add_argument('--start-timeout', type=float, default=600.0, help='seconds for the PE to start run.sh')
    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('--stall-reset', type=float, default=360.0, help='reset the VM if a dark screen is frozen this long while disk and serial are idle')
    p.add_argument('--max-resets', type=int, default=6)
    p.add_argument('--reboot-timeout', type=float, default=60.0, help='seconds after a guest reboot request without a new kernel boot before the VM is reset')
    p.add_argument('--halt-timeout', type=float, default=150.0, help='a pure-black, idle screen this long means the guest halted')
    p.add_argument('--progress-file', default=None, help='the system disk; its mtime shows guest disk activity')
    p.add_argument('--disk-idle', type=float, default=90.0, help='seconds without disk/serial activity that count as idle')
    p.add_argument('--bright', type=float, default=80.0, help='mean brightness above which a screen is the loginwindow/desktop')
    p.add_argument('--login-grace', type=float, default=180.0, help='seconds a bright screen must persist before powering down')
    p.add_argument('qemu', nargs=argparse.REMAINDER)
    args = p.parse_args()
    qemu_args = [a for a in args.qemu if a != '--']
    if not qemu_args:
        p.error('QEMU command line required after --')

    debug_dir = args.debug_dir or f'/tmp/vmix-macos/{args.name}'
    debug_dir = prepare_debug_dir(debug_dir)
    log = Log(os.path.join(debug_dir, 'driver.log'))
    log(f'mode={args.mode} debug-dir={debug_dir} serial={args.serial_log}')

    qmp_sock = os.path.join(debug_dir, 'qmp.sock')
    proc = launch(qemu_args, qmp_sock, log)
    qmp = QMP(qmp_sock)
    try:
        qmp.connect()
    except Exception as e:  # noqa: BLE001
        log(f'QMP connect failed: {e}')
        proc.kill()
        return 2
    screen = Screen(qmp, debug_dir, log)
    serial = Serial(args.serial_log)
    try:
        if args.mode == 'qga':
            rc = run_qga(args, proc, qmp, screen, serial, log)
        else:
            rc = drive(args, proc, qmp, screen, serial, log)
    finally:
        if args.serial_log and os.path.exists(args.serial_log):
            try:
                shutil.copy(args.serial_log, os.path.join(debug_dir, 'serial.log'))
            except OSError:
                pass
        try:
            if proc.poll() is None:
                proc.kill()
        except Exception:  # noqa: BLE001
            pass
    log(f'done rc={rc}')
    return rc


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