#!/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


state_dir = Path(env("WIN11_VM_STATE_DIR") or script_dir / "vm" / "windows11")
disk_size = env("WIN11_VM_DISK_SIZE") or "64G"
rdp_port = env("WIN11_VM_RDP_PORT")
ssh_port = env("WIN11_VM_SSH_PORT")
forwards_file = Path(
    env("WIN11_VM_FORWARDS") or script_dir / "qemu" / "windows11" / "forwards.conf")
qemu_config_template = Path(
    env("WIN11_VM_QEMU_CONFIG") or script_dir / "qemu" / "windows11" / "config.ini")


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


daemonize = False
install_from_iso = False
install_from_img = False
download_only = False
seed_only = False
extra_args = []
arch_option = None
vnc_listen = None


def usage(stream=sys.stdout):
    stream.write(f"""\
Usage: {Path(sys.argv[0]).name} [--daemon] [--download-only] [--seed-only] [--vnc[=ADDR:N]] [--install-from-iso | --install-from-img]
       [--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 from
the Windows 11 QCOW2 base image.
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.
Use --install-from-iso to perform an automatic ISO installation instead, or
--install-from-img to install from a disk image built from the ISO.
If it exists, start Windows normally without attaching installation media.

Options:
  --daemon            Run QEMU in the background
  --download-only     Download the VirtIO tools ISO and the Windows 11 ISO
                      (for --arch) if they are missing, then exit; needs
                      Node.js and Chromium for the latter, see vm-images/README.md
                      (the base QCOW2 is supplied manually)
  --seed-only         Build the seed ISO (answer file and setup scripts) and exit,
                      without starting the VM or fetching the Windows and VirtIO ISOs
  --install-from-iso  Install from ISO when the system disk is missing
  --install-from-img  Like --install-from-iso, but install from the disk image
                      made by vm-images/build-windows-install-disk.sh, which
                      is the ISO path ending in .img instead of .iso.
                      It is attached as an NVMe disk, much faster than the
                      emulated USB CD-ROM, and is never modified
  --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, seed ISO 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_INSTALL_ISO  Windows 11 installation ISO used with --install-from-iso
                        (--install-from-img uses the same path ending in .img)
  WIN11_VM_TOOLS_ISO    VirtIO Windows tools ISO (default: resolved by shared image helper)
  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     Answer-file 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)
""")


args = sys.argv[1:]
while args:
    arg = args.pop(0)
    if arg == "--daemon":
        daemonize = True
    elif arg == "--download-only":
        download_only = True
    elif arg == "--seed-only":
        seed_only = True
    elif arg == "--vnc" or arg.startswith("--vnc="):
        vnc_listen = arg.partition("=")[2] or "127.0.0.1:0"
        if ":" not in vnc_listen:
            vnc_listen += ":0"
    elif arg == "--install-from-iso":
        install_from_iso = True
    elif arg == "--install-from-img":
        install_from_img = True
    elif arg == "--state-dir":
        if not args:
            die(f"{arg} requires a directory argument.", 2)
        state_dir = Path(args.pop(0))
    elif arg.startswith("--state-dir="):
        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":
            rdp_port = value
        else:
            ssh_port = value
    elif arg == "--arch":
        if not args:
            die(f"{arg} requires an architecture argument.", 2)
        arch_option = args.pop(0)
    elif arg.startswith("--arch="):
        arch_option = arg.split("=", 1)[1]
    elif arg in ("-h", "--help"):
        usage()
        sys.exit(0)
    elif arg == "--":
        extra_args.extend(args)
        break
    else:
        print(f"Error: unknown option: {arg}", file=sys.stderr)
        usage(sys.stderr)
        sys.exit(2)

host_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"] = host_arch
qemu_binary = "qemu-system-aarch64" if host_arch == "arm64" else "qemu-system-x86_64"
host_os = platform.system().lower()
if host_os == "darwin" and host_arch == "amd64":
    die("amd64 guests are not supported on macOS hosts; use an arm64 image")
qemu_platform_config = script_dir / "qemu" / "windows11" / f"{host_os}-{host_arch}.ini"


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


if host_arch == "arm64":
    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",
    )
    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:
    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",
    )
    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",
    )
if install_from_img:
    install_from_iso = True

if download_only:
    # The VirtIO tools ISO and the Windows installation ISO are the media the
    # launcher downloads itself; the base QCOW2 is supplied manually.
    image_downloader = None
    if env("WIN11_VM_TOOLS_ISO"):
        tools_iso = Path(env("WIN11_VM_TOOLS_ISO"))
        if not tools_iso.is_file():
            die(f"VirtIO Windows tools ISO was not found: {tools_iso}")
    else:
        image_downloader = find_image_downloader()
        tools_iso = Path(capture([str(image_downloader), "--print-path", "virtio-win"]))
        if not tools_iso.is_file():
            run([str(image_downloader), "virtio-win"])
        if not tools_iso.is_file():
            die(f"downloader did not create: {tools_iso}")
    print(f"VirtIO Windows tools ISO is ready: {tools_iso}")
    if env("WIN11_VM_INSTALL_ISO"):
        install_iso = Path(env("WIN11_VM_INSTALL_ISO"))
        if not install_iso.is_file():
            die(f"Windows 11 installation ISO was not found: {install_iso}")
    else:
        image_downloader = image_downloader or find_image_downloader()
        run([str(image_downloader), "windows11"])
        install_iso = Path(capture([str(image_downloader), "--print-path", "windows11", "--iso"]))
    print(f"Windows 11 installation ISO is ready: {install_iso}")
    sys.exit(0)

disk_image = Path(env("WIN11_VM_DISK_IMAGE") or state_dir / "disk.qcow2")
seed_iso = Path(env("WIN11_VM_SEED_ISO") or disk_image.parent / "seed.iso")
uefi_vars = state_dir / "OVMF_VARS_4M.fd"
pid_file = state_dir / "qemu.pid"
log_file = state_dir / "console.log"
monitor_socket = state_dir / "qemu-monitor.sock"
guest_agent_socket = state_dir / "qemu-guest-agent.sock"
qemu_config = state_dir / "qemu.ini"


def build_seed_iso(image_downloader, force=False):
    """Build the seed ISO, unless it is newer than all of its sources and not forced."""
    seed_dir = script_dir / "seediso" / "windows11"
    seed_sources = [seed_dir / "Autounattend.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.
    seed_sources.append(Path(capture(
        [str(image_downloader), "--print-path", "openssh-win"])))

    def newer_than_seed(source):
        # Mirrors bash `-nt`: true if source exists and is newer than seed_iso.
        return source.exists() and source.stat().st_mtime > seed_iso.stat().st_mtime

    if force or not seed_iso.is_file() or not seed_sources[-1].is_file() \
            or any(newer_than_seed(s) for s in seed_sources):
        run([str(script_dir / "seediso" / "windows11" / "build-seed.sh"), str(seed_iso), host_arch])


if seed_only:
    build_seed_iso(find_image_downloader(), force=True)
    print(f"Seed ISO is ready: {seed_iso}")
    sys.exit(0)

installing = False
installing_from_iso = False
if not disk_image.is_file():
    installing = True
    installing_from_iso = install_from_iso
    image_downloader = find_image_downloader()
    if install_from_iso:
        install_iso = Path(
            env("WIN11_VM_INSTALL_ISO")
            or capture([str(image_downloader), "--print-path", "windows11", "--iso"]))
        if install_from_img:
            install_iso = install_iso.with_suffix(".img")
        uses_shared_tools_iso = False
        if env("WIN11_VM_TOOLS_ISO"):
            tools_iso = Path(env("WIN11_VM_TOOLS_ISO"))
        else:
            tools_iso = Path(capture([str(image_downloader), "--print-path", "virtio-win"]))
            uses_shared_tools_iso = True
    else:
        base_image = Path(env("WIN11_VM_BASE_IMAGE")
                          or capture([str(image_downloader), "--print-path", "windows11"]))

if shutil.which(qemu_binary) is None:
    die(f"{qemu_binary} is not installed.")
if installing:
    if shutil.which("qemu-img") is None:
        die("qemu-img is not installed.")
    if installing_from_iso:
        if not install_iso.is_file():
            if install_from_img:
                die(f"Windows 11 installation disk image was not found: {install_iso}\n"
                    "Build it with vm-images/build-windows-install-disk.sh ISO")
            die(f"Windows 11 ISO was not found: {install_iso}")
    elif not base_image.is_file():
        die(f"Windows 11 base QCOW2 image was not found: {base_image}")
for firmware_file in (uefi_code, uefi_vars_template):
    if not firmware_file.is_file():
        die(f"required UEFI firmware was not found: {firmware_file}")

state_dir.mkdir(parents=True, exist_ok=True)
if installing_from_iso:
    if not tools_iso.is_file():
        if uses_shared_tools_iso:
            run([str(image_downloader), "virtio-win"])
        else:
            die(f"VirtIO Windows tools ISO was not found: {tools_iso}")
    if not tools_iso.is_file():
        die(f"downloader did not create: {tools_iso}")

    build_seed_iso(image_downloader)

if pid_file.is_file():
    old_pid = 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}).")
    pid_file.unlink(missing_ok=True)
monitor_socket.unlink(missing_ok=True)
guest_agent_socket.unlink(missing_ok=True)
if not uefi_vars.is_file():
    shutil.copyfile(uefi_vars_template, uefi_vars)
if installing:
    if installing_from_iso:
        run(["qemu-img", "create", "-q", "-f", "qcow2", str(disk_image), disk_size])
    else:
        run(["qemu-img", "create", "-q", "-f", "qcow2", "-F", "qcow2",
             "-b", str(base_image), str(disk_image), disk_size])
    print(f"Created Windows 11 system disk: {disk_image}")

# 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.
display_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)}")

# Most of the VM definition lives in the QEMU config template; fill in the
# file paths and hand the rendered file to QEMU with -readconfig.
config_values = {
    "uefi_code": uefi_code,
    "uefi_vars": uefi_vars,
    "disk_image": disk_image,
    "guest_agent_socket": guest_agent_socket,
    "monitor_socket": 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 (qemu_config_template, qemu_platform_config):
    if not config_file.is_file():
        die(f"QEMU config file was not found: {config_file}")
try:
    qemu_config_text = string.Template(qemu_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 {qemu_config_template}: {error}")
qemu_config.write_text(qemu_config_text)

# Each rule of the port forward file is one hostfwd entry of the network.
forwards = vmforwards.usable_forwards(vmforwards.load_forwards(
    forwards_file, {label: port for label, port in (("RDP", rdp_port), ("SSH", ssh_port)) if port}))

# Neither Windows Setup nor the base image has the viogpudo driver yet, and
# without it the virtio-gpu in the config file shows "display output is not
# active". ramfb shows the screen until Configure-Windows.ps1 installs it. It
# must be created before the virtio-gpu: the display window shows the first GPU.
qemu_args = ["-device", "ramfb"] if installing and host_arch == "arm64" else []
qemu_args += [
    "-readconfig", str(qemu_config),
    "-readconfig", str(qemu_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,
    # "-display", "none", "-spice", "port=5900,addr=127.0.0.1,disable-ticketing=on,agent-mouse=off", "-device", "virtio-vga",
]
if host_os == "darwin" and host_arch == "arm64" and not install_from_img:
    # With the HVF in-kernel GICv3, QEMU's usb-storage races on the status
    # packet and Windows Setup stalls 10 s after every burst of reads
    # (utmapp/UTM#7970), so use the emulated interrupt controller; upstream
    # QEMU only gets the fix from the UTM fork. The NVMe install disk of
    # --install-from-img does not use usb-storage, and hangs with the emulated
    # controller.
    qemu_args += ["-machine", "kernel-irqchip=off"]
if installing_from_iso:
    # q35 has IDE for the installation media; "virt" does not, so use USB.
    for media_id, media_file, boot_index, ide_bus in (
            ("install_media", install_iso, 2, "ide.1"),
            ("answer_media", seed_iso, 3, "ide.2"),
            ("tools_media", tools_iso, 4, "ide.3")):
        device = (f"ide-cd,drive={media_id},bus={ide_bus},bootindex={boot_index}"
                  if host_arch == "amd64"
                  else f"usb-storage,drive={media_id},bootindex={boot_index}")
        if media_id == "install_media" and install_from_img:
            # snapshot=on: NVMe cannot report write protection, and Windows
            # writes to the volume it mounts, so a read-only drive makes it fail.
            # The writes go to a temporary overlay and the image stays unchanged.
            qemu_args += [
                "-drive", f"file={media_file},if=none,id={media_id},format=raw,snapshot=on",
                "-device", f"nvme,drive={media_id},serial=installer,bootindex={boot_index}",
            ]
            continue
        qemu_args += [
            "-drive", f"file={media_file},if=none,id={media_id},media=cdrom,readonly=on",
            "-device", device,
        ]

if installing_from_iso:
    print("Starting Windows 11 automatic installation.")
elif installing:
    print("Starting Windows 11 from the QCOW2 base image.")
else:
    print("Starting Windows 11 normally.")
print("Account: guofang / password; " + "; ".join(
    f"{forward.label}: 0.0.0.0:{forward.host_port}" for forward in forwards),
    flush=True)

if installing_from_iso and not install_from_img:
    # The Windows ISO boots through bootmgr, which waits for a key ("Press any
    # key to boot from CD or DVD...") and gives up after a few seconds. Nobody
    # is there to press it, so a detached helper taps the space bar through the
    # QEMU monitor for the first 30 s. It exits by itself; later key presses
    # are harmless because Setup is unattended.
    key_script = (
        "import socket, sys, time\n"
        "path = sys.argv[1]\n"
        "end = time.time() + 30\n"
        "while time.time() < end:\n"
        "    try:\n"
        "        s = socket.socket(socket.AF_UNIX)\n"
        "        s.connect(path)\n"
        "        break\n"
        "    except OSError:\n"
        "        time.sleep(0.5)\n"
        "else:\n"
        "    sys.exit(0)\n"
        "while time.time() < end:\n"
        "    s.sendall(b'sendkey spc\\n')\n"
        "    time.sleep(1)\n")
    subprocess.Popen([sys.executable, "-c", key_script, str(monitor_socket)],
                     stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
                     stderr=subprocess.DEVNULL, start_new_session=True)

if daemonize:
    log_file.write_text("")
    qemu_args += ["-serial", f"file:{log_file}", "-daemonize", "-pidfile", str(pid_file)]
    run([qemu_binary, *qemu_args, *extra_args])
    print(f"VM started (PID {pid_file.read_text().strip()}); console log: {log_file}")
else:
    os.execvp(qemu_binary, [qemu_binary, *qemu_args, *extra_args])