#!/usr/bin/env python3
"""Start the Windows 11 virtual machine."""

import os
import platform
import shutil
import string
import subprocess
import sys
from pathlib import Path

import vmforwards

script_dir = Path(__file__).resolve().parent
env = os.environ.get


def die(message, code=1):
    print(f"Error: {message}", file=sys.stderr)
    sys.exit(code)


def run(command):
    """Run a command, exiting with its status on failure (like `set -e`)."""
    result = subprocess.run(command)
    if result.returncode != 0:
        sys.exit(result.returncode)


def capture(command):
    """Return a command's stdout without the trailing newline, like $(...)."""
    result = subprocess.run(command, stdout=subprocess.PIPE, text=True)
    if result.returncode != 0:
        sys.exit(result.returncode)
    return result.stdout.rstrip("\n")


def is_executable(path):
    return path.is_file() and os.access(path, os.X_OK)


def find_image_downloader():
    if env("VM_IMAGE_HELPER"):
        image_downloader = Path(env("VM_IMAGE_HELPER"))
    elif is_executable(script_dir / "vm-images" / "download-images.sh"):
        image_downloader = script_dir / "vm-images" / "download-images.sh"
    else:
        image_downloader = script_dir / ".." / "vm-images" / "download-images.sh"
    if not is_executable(image_downloader):
        die(f"shared image helper is unavailable: {image_downloader}")
    return image_downloader


def normalize_arch(arch):
    if arch in ("arm64", "aarch64"):
        return "arm64"
    if arch in ("x86_64", "amd64"):
        return "amd64"
    die(f"unsupported architecture: {arch}")


def pick_firmware(override, *candidates):
    if override:
        return Path(override)
    for candidate in candidates:
        if Path(candidate).is_file():
            return Path(candidate)
    return Path(candidates[0])


class Options:
    """Command-line options."""

    def __init__(self):
        self.daemonize = False
        self.vnc_listen = None
        self.state_dir = Path(env("WIN11_VM_STATE_DIR") or script_dir / "vm" / "windows11")
        self.rdp_port = env("WIN11_VM_RDP_PORT")
        self.ssh_port = env("WIN11_VM_SSH_PORT")
        self.arch_option = None
        self.extra_args = []


def usage(stream=sys.stdout):
    stream.write(f"""\
Usage: {Path(sys.argv[0]).name} [--daemon] [--vnc[=ADDR:N]]
       [--arch ARCH] [--state-dir DIR] [--rdp-port PORT] [--ssh-port PORT] [-- QEMU_ARGUMENT ...]

If the system disk (DIR/disk.qcow2, see --state-dir) is missing, create it as a
QCOW2 overlay on the Windows 11 base image, which starts in seconds and runs the
specialize and OOBE passes of its answer file on the first boot, and mount the
seed ISO (DIR/seed.iso, built from seediso/windows11 if missing or outdated):
Windows reads its Unattend.xml from the root of that CD-ROM during the specialize
and OOBE passes: host name, account, auto logon, Remote Desktop, then at the
first logon Configure-Windows.ps1 (OpenSSH and the SSH key, settings, installer
scripts on the desktop). If the disk exists, start Windows normally, without the
seed ISO.
Make the base image once with vm-images/build-windows-base.sh; it installs
Windows from the ISO and generalizes the result. The base image contains only
Windows (no account, no SSH, no VirtIO or other hardware-specific drivers), so
the devices in qemu/windows11 are ones Windows has inbox drivers for.
The VM hardware (accelerator, memory, CPUs, drives, devices) is defined in
qemu/windows11/config.ini, which is rendered to qemu.ini in the state directory and
loaded with QEMU's -readconfig, followed by qemu/windows11/<os>-<arch>.ini (os is linux or darwin, arch is amd64 or arm64, chosen from the host system and
architecture, --arch or VM_IMAGE_ARCH). Later QEMU arguments override them, e.g. -- -m 4G -smp 4.

Options:
  --daemon            Run QEMU in the background
  --vnc[=ADDR:N]      Serve the screen over VNC instead of opening a window, on
                      ADDR (default: 127.0.0.1) at TCP port 5900+N (default N: 0).
                      Overrides WIN11_VM_DISPLAY
  --arch ARCH         Guest architecture: arm64 or amd64 (default: the host's;
                      overrides VM_IMAGE_ARCH). A guest that differs from the
                      host runs under slow TCG emulation
  --state-dir DIR     Directory for the VM disk and runtime files
                      (default: vm/windows11; overrides WIN11_VM_STATE_DIR)
  --rdp-port PORT     Host port of the RDP rule in the port forward file
                      (default: 3389; overrides WIN11_VM_RDP_PORT)
  --ssh-port PORT     Host port of the SSH rule in the port forward file
                      (default: 2222; overrides WIN11_VM_SSH_PORT)
  -h, --help          Show this help

Environment variables:
  WIN11_VM_BASE_IMAGE   Windows 11 base QCOW2 image (default: beside shared images)
  WIN11_VM_RDP_PORT     Host RDP port (default: from the port forward file, 3389)
  WIN11_VM_SSH_PORT     Host SSH port (default: from the port forward file, 2222)
  WIN11_VM_FORWARDS     Port forward file (default: qemu/windows11/forwards.conf)
  WIN11_VM_STATE_DIR    Runtime directory (default: vm/windows11)
  WIN11_VM_DISK_IMAGE   System disk path
  WIN11_VM_DISK_SIZE    New system disk size (default: 64G)
  WIN11_VM_SEED_ISO     Seed ISO path (default: alongside the VM disk)
  WIN11_VM_UEFI_CODE    OVMF read-only firmware path (auto-detected if unset)
  WIN11_VM_UEFI_VARS    OVMF variable-store template path (auto-detected if unset)
  WIN11_VM_QEMU_CONFIG  QEMU config template passed to -readconfig
                        (default: qemu/windows11/config.ini)
  WIN11_VM_DISPLAY      QEMU display backend (default: cocoa on macOS,
                        gtk elsewhere)
""")


def parse_args(args):
    """Parse the command line."""
    options = Options()
    args = list(args)
    while args:
        arg = args.pop(0)
        if arg == "--daemon":
            options.daemonize = True
        elif arg == "--vnc" or arg.startswith("--vnc="):
            options.vnc_listen = arg.partition("=")[2] or "127.0.0.1:0"
            if ":" not in options.vnc_listen:
                options.vnc_listen += ":0"
        elif arg == "--state-dir":
            if not args:
                die(f"{arg} requires a directory argument.", 2)
            options.state_dir = Path(args.pop(0))
        elif arg.startswith("--state-dir="):
            options.state_dir = Path(arg.split("=", 1)[1])
        elif arg in ("--rdp-port", "--ssh-port") or arg.startswith(("--rdp-port=", "--ssh-port=")):
            name, sep, value = arg.partition("=")
            if not sep:
                if not args:
                    die(f"{name} requires a port argument.", 2)
                value = args.pop(0)
            if not (value.isascii() and value.isdigit() and 1 <= int(value) <= 65535):
                die(f"{name} must be a port number from 1 to 65535: {value}", 2)
            if name == "--rdp-port":
                options.rdp_port = value
            else:
                options.ssh_port = value
        elif arg == "--arch":
            if not args:
                die(f"{arg} requires an architecture argument.", 2)
            options.arch_option = args.pop(0)
        elif arg.startswith("--arch="):
            options.arch_option = arg.split("=", 1)[1]
        elif arg in ("-h", "--help"):
            usage()
            sys.exit(0)
        elif arg == "--":
            options.extra_args.extend(args)
            break
        else:
            print(f"Error: unknown option: {arg}", file=sys.stderr)
            usage(sys.stderr)
            sys.exit(2)
    return options


class Platform:
    """Guest architecture, host system and the UEFI firmware that fit them."""

    def __init__(self, arch_option):
        self.arch = normalize_arch(arch_option or env("VM_IMAGE_ARCH") or platform.machine())
        # The shared image helper and build-seed.sh pick the architecture from here.
        os.environ["VM_IMAGE_ARCH"] = self.arch
        self.qemu_binary = "qemu-system-aarch64" if self.arch == "arm64" else "qemu-system-x86_64"
        self.os = platform.system().lower()
        if self.os == "darwin" and self.arch == "amd64":
            die("amd64 guests are not supported on macOS hosts; use an arm64 image")
        self.config_template = Path(
            env("WIN11_VM_QEMU_CONFIG") or script_dir / "qemu" / "windows11" / "config.ini")
        self.platform_config = script_dir / "qemu" / "windows11" / f"{self.os}-{self.arch}.ini"
        if self.arch == "arm64":
            self.uefi_code = pick_firmware(
                env("WIN11_VM_UEFI_CODE"),
                "/usr/share/AAVMF/AAVMF_CODE.fd",
                "/usr/share/edk2/aarch64/QEMU_EFI.fd",
                "/opt/homebrew/share/qemu/edk2-aarch64-code.fd",
                "/usr/local/share/qemu/edk2-aarch64-code.fd",
            )
            self.uefi_vars_template = pick_firmware(
                env("WIN11_VM_UEFI_VARS"),
                "/usr/share/AAVMF/AAVMF_VARS.fd",
                "/usr/share/edk2/aarch64/vars-template-pflash.raw",
                "/opt/homebrew/share/qemu/edk2-arm-vars.fd",
                "/usr/local/share/qemu/edk2-arm-vars.fd",
            )
        else:
            self.uefi_code = pick_firmware(
                env("WIN11_VM_UEFI_CODE"),
                "/usr/share/OVMF/OVMF_CODE_4M.fd",
                "/usr/share/edk2/ovmf/OVMF_CODE.fd",
                "/opt/homebrew/share/qemu/edk2-x86_64-code.fd",
                "/usr/local/share/qemu/edk2-x86_64-code.fd",
            )
            self.uefi_vars_template = pick_firmware(
                env("WIN11_VM_UEFI_VARS"),
                "/usr/share/OVMF/OVMF_VARS_4M.fd",
                "/usr/share/edk2/ovmf/OVMF_VARS.fd",
                "/opt/homebrew/share/qemu/edk2-i386-vars.fd",
                "/usr/local/share/qemu/edk2-i386-vars.fd",
            )


class Runtime:
    """Paths of the files a VM keeps in its state directory."""

    def __init__(self, state_dir):
        self.state_dir = state_dir
        self.disk_image = Path(env("WIN11_VM_DISK_IMAGE") or state_dir / "disk.qcow2")
        self.uefi_vars = state_dir / "OVMF_VARS_4M.fd"
        self.pid_file = state_dir / "qemu.pid"
        self.log_file = state_dir / "console.log"
        self.monitor_socket = state_dir / "qemu-monitor.sock"
        self.guest_agent_socket = state_dir / "qemu-guest-agent.sock"
        self.qemu_config = state_dir / "qemu.ini"


def check_tools_and_firmware(plat, need_qemu_img=True):
    if shutil.which(plat.qemu_binary) is None:
        die(f"{plat.qemu_binary} is not installed.")
    if need_qemu_img and shutil.which("qemu-img") is None:
        die("qemu-img is not installed.")
    for firmware_file in (plat.uefi_code, plat.uefi_vars_template):
        if not firmware_file.is_file():
            die(f"required UEFI firmware was not found: {firmware_file}")


def prepare_runtime(plat, rt):
    """Create the state directory and clear the leftovers of an earlier run."""
    rt.state_dir.mkdir(parents=True, exist_ok=True)
    if rt.pid_file.is_file():
        old_pid = rt.pid_file.read_text().strip()
        if old_pid.isdigit():
            try:
                os.kill(int(old_pid), 0)
            except OSError:
                pass
            else:
                die(f"Windows 11 VM is already running (PID {old_pid}).")
        rt.pid_file.unlink(missing_ok=True)
    rt.monitor_socket.unlink(missing_ok=True)
    rt.guest_agent_socket.unlink(missing_ok=True)
    if not rt.uefi_vars.is_file():
        shutil.copyfile(plat.uefi_vars_template, rt.uefi_vars)


def display_backend(vnc_listen):
    """cocoa shows the guest at its own resolution. Append ,zoom-to-fit=on (via
    WIN11_VM_DISPLAY) to stretch a low-resolution screen to the window."""
    backend = (f"vnc={vnc_listen}" if vnc_listen else env("WIN11_VM_DISPLAY") or (
        "cocoa" if platform.system() == "Darwin" else "gtk"))
    if vnc_listen:
        vnc_host, _, vnc_display = vnc_listen.rpartition(":")
        if not vnc_display.isdigit():
            die(f"--vnc display number must be an integer: {vnc_listen}", 2)
        print(f"VNC: connect to {vnc_host or '127.0.0.1'}:{5900 + int(vnc_display)}")
    return backend


def common_qemu_args(plat, rt, options):
    """Render the config template and return the QEMU arguments shared by both
    launchers, plus the port forwards (for the start-up message).

    Most of the VM definition lives in the QEMU config template; the file
    paths are filled in and the rendered file is handed to QEMU with -readconfig.
    """
    config_values = {
        "uefi_code": plat.uefi_code,
        "uefi_vars": rt.uefi_vars,
        "disk_image": rt.disk_image,
        "guest_agent_socket": rt.guest_agent_socket,
        "monitor_socket": rt.monitor_socket,
    }
    for key, value in config_values.items():
        # QEMU's config file parser has no escape syntax for quoted values.
        if any(char in str(value) for char in '"\n'):
            die(f"{key} must not contain double quotes or newlines: {value}")
    for config_file in (plat.config_template, plat.platform_config):
        if not config_file.is_file():
            die(f"QEMU config file was not found: {config_file}")
    try:
        config_text = string.Template(plat.config_template.read_text()).substitute(
            {key: str(value) for key, value in config_values.items()})
    except (KeyError, ValueError) as error:
        die(f"invalid placeholder in {plat.config_template}: {error}")
    rt.qemu_config.write_text(config_text)

    # Each rule of the port forward file is one hostfwd entry of the network.
    forwards_file = Path(
        env("WIN11_VM_FORWARDS") or script_dir / "qemu" / "windows11" / "forwards.conf")
    forwards = vmforwards.usable_forwards(vmforwards.load_forwards(
        forwards_file,
        {label: port for label, port in (("RDP", options.rdp_port), ("SSH", options.ssh_port))
         if port}))

    qemu_args = [
        "-readconfig", str(rt.qemu_config),
        "-readconfig", str(plat.platform_config),
        # "max" is the host CPU under KVM/HVF and the fullest model under TCG.
        "-cpu", "max",
        # A config file section cannot hold two hostfwd entries, so the network
        # that the NIC in the config file uses is defined here. -nic none stops
        # QEMU from adding its default NIC as well.
        "-netdev", vmforwards.netdev_argument(forwards),
        "-nic", "none",
        "-display", display_backend(options.vnc_listen),
    ]
    return qemu_args, forwards


def start_qemu(plat, rt, qemu_args, extra_args, forwards, daemonize, banner):
    print(banner)
    print("Account: guofang / password; " + "; ".join(
        f"{forward.label}: 0.0.0.0:{forward.host_port}" for forward in forwards),
        flush=True)
    if daemonize:
        rt.log_file.write_text("")
        qemu_args += ["-serial", f"file:{rt.log_file}", "-daemonize", "-pidfile", str(rt.pid_file)]
        run([plat.qemu_binary, *qemu_args, *extra_args])
        print(f"VM started (PID {rt.pid_file.read_text().strip()}); console log: {rt.log_file}")
    else:
        os.execvp(plat.qemu_binary, [plat.qemu_binary, *qemu_args, *extra_args])


def build_seed_iso(seed_iso):
    """Build the seed ISO, unless it is newer than all of its sources."""
    seed_dir = script_dir / "seediso" / "windows11"
    seed_sources = [seed_dir / "Unattend.xml", seed_dir / "Configure-Windows.ps1",
                    *sorted(seed_dir.glob("[0-9][0-9]-*.bat")),
                    *sorted(seed_dir.glob("[0-9][0-9]-*.ps1"))]

    # build-seed.sh embeds the administrator SSH public key.
    ssh_key = os.environ.get("SSH_PUBLIC_KEY_FILE") or next(
        (str(k) for n in ("id_ed25519", "id_ecdsa", "id_rsa")
         if (k := Path.home() / ".ssh" / f"{n}.pub").is_file()), None)
    if ssh_key:
        seed_sources.append(Path(ssh_key))

    # The seed carries the OpenSSH zip for this architecture; build-seed.sh
    # downloads it when missing.
    openssh_zip = Path(capture(
        [str(find_image_downloader()), "--print-path", "openssh-win"]))
    seed_sources.append(openssh_zip)

    if not seed_iso.is_file() or not openssh_zip.is_file() or any(
            source.exists() and source.stat().st_mtime > seed_iso.stat().st_mtime
            for source in seed_sources):
        run([str(seed_dir / "build-seed.sh"), str(seed_iso), plat.arch])


options = parse_args(sys.argv[1:])
plat = Platform(options.arch_option)
rt = Runtime(options.state_dir)
disk_size = env("WIN11_VM_DISK_SIZE") or "64G"

creating = not rt.disk_image.is_file()
if creating:
    base_image = Path(env("WIN11_VM_BASE_IMAGE")
                      or capture([str(find_image_downloader()), "--print-path", "windows11"]))
    if not base_image.is_file():
        die(f"Windows 11 base QCOW2 image was not found: {base_image}\n"
            "Make it once with vm-images/build-windows-base.sh")

check_tools_and_firmware(plat, need_qemu_img=creating)
prepare_runtime(plat, rt)
if creating:
    seed_iso = Path(env("WIN11_VM_SEED_ISO") or rt.disk_image.parent / "seed.iso")
    build_seed_iso(seed_iso)
if creating:
    run(["qemu-img", "create", "-q", "-f", "qcow2", "-F", "qcow2",
         "-b", str(base_image), str(rt.disk_image), disk_size])
    print(f"Created Windows 11 system disk: {rt.disk_image}")

qemu_args, forwards = common_qemu_args(plat, rt, options)
if plat.os == "darwin" and plat.arch == "arm64":
    # The HVF in-kernel GICv3 stalls Windows on its USB and NVMe devices
    # (utmapp/UTM#7970), so use the emulated interrupt controller.
    qemu_args += ["-machine", "kernel-irqchip=off"]
if creating:
    # q35 has IDE (ide.0 is the system disk); "virt" does not, so use USB.
    qemu_args += ["-drive", f"file={seed_iso},if=none,id=seed_media,media=cdrom,readonly=on",
                  "-device", "ide-cd,drive=seed_media,bus=ide.1" if plat.arch == "amd64"
                  else "usb-storage,drive=seed_media"]
start_qemu(plat, rt, qemu_args, options.extra_args, forwards, options.daemonize,
           "Starting Windows 11 from the QCOW2 base image." if creating
           else "Starting Windows 11 normally.")