#!/usr/bin/env bash
# Build the Windows 11 base QCOW2 image that run-windows11.py clones.
#
# run-windows11.py creates a missing system disk as a QCOW2 overlay on this
# base image, so a new VM starts in seconds instead of running Windows Setup
# again. The base image is made once:
#   1. run-windows11.py --install-from-iso installs Windows unattended into a
#      scratch state directory (Autounattend.xml, Configure-Windows.ps1: VirtIO
#      drivers, OpenSSH, the administrator key from build-seed.sh);
#   2. once SSH answers, this script runs sysprep /generalize over SSH, which
#      removes the machine identity and shuts the guest down;
#   3. the shut-down disk is converted into the output file and made read-only.
# The first boot of every clone then runs the specialize and OOBE passes of the
# answer file below (host name, account, automatic logon) without any prompt.
#
# Needs the Windows ISO in install-iso/ (see README.md), qemu-img, ssh, scp and
# a private key that matches the public key put on the answer media (the
# default ~/.ssh keys, or SSH_PUBLIC_KEY_FILE for the public half and
# SSH_PRIVATE_KEY_FILE for this script).
#
# Usage: build-windows-base.sh [--force] [--no-sysprep] [--keep-work] [--compress]
#                              [--window] [--ssh-port PORT] [OUTPUT.qcow2]
# OUTPUT defaults to the path that "download-images.sh --print-path windows11"
# prints, which is where run-windows11.py looks for the base image.
set -euo pipefail

images_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
repo_dir="$(cd -- "${images_dir}/.." && pwd)"
launcher="${repo_dir}/run-windows11.py"
downloader="${images_dir}/download-images.sh"

usage() {
  cat <<EOF
Usage: $(basename "$0") [--force] [--no-sysprep] [--keep-work] [--compress]
                         [--window] [--ssh-port PORT] [OUTPUT.qcow2]

Install Windows 11 from the ISO, generalize it with sysprep and store the disk
as the read-only base image that run-windows11.py clones.

  --force         Replace OUTPUT if it exists
  --no-sysprep    Skip sysprep and only shut the guest down; clones then share
                  the machine SID and host name (fine for stand-alone VMs)
  --keep-work     Keep the scratch state directory (also kept on failure)
  --compress      Compress the output (smaller, but clones read it slower)
  --window        Show the guest in a QEMU window (needs a desktop session)
                  instead of the VNC display on 127.0.0.1:5990
  --ssh-port PORT Host port forwarded to the guest's SSH during the build
                  (default: 2299; overrides BASE_BUILD_SSH_PORT)
  -h, --help      Show this help

Environment variables:
  SSH_PRIVATE_KEY_FILE  Private key to log in with (default: ssh's own choice)
  SSH_PUBLIC_KEY_FILE   Public key put on the answer media (see build-seed.sh)
  VM_IMAGE_ARCH         amd64 or arm64 (default: the host's)
  BASE_BUILD_TMPDIR     Parent of the scratch state directory, which needs about
                        25 GB (default: /var/tmp)
  BASE_BUILD_TIMEOUT    Seconds to wait for the installation (default: 5400)
EOF
}

die() {
  printf 'Error: %s\n' "$1" >&2
  exit 1
}

force=0
sysprep=1
keep_work=0
compress=0
window=0
ssh_port="${BASE_BUILD_SSH_PORT:-2299}"
output=""
while (($#)); do
  case "$1" in
    -h|--help) usage; exit 0 ;;
    --force) force=1 ;;
    --no-sysprep) sysprep=0 ;;
    --keep-work) keep_work=1 ;;
    --compress) compress=1 ;;
    --window) window=1 ;;
    --ssh-port) shift; ssh_port="${1:?--ssh-port needs a value}" ;;
    --ssh-port=*) ssh_port="${1#*=}" ;;
    -*) usage >&2; exit 2 ;;
    *) [[ -z "$output" ]] || { usage >&2; exit 2; }; output="$1" ;;
  esac
  shift
done
[[ "$ssh_port" =~ ^[0-9]+$ ]] || die "--ssh-port must be a number: ${ssh_port}"

for tool in qemu-img ssh scp python3; do
  command -v "$tool" >/dev/null || die "${tool} is not installed."
done
[[ -x "$launcher" ]] || die "run-windows11.py was not found: ${launcher}"
[[ -x "$downloader" ]] || die "download-images.sh was not found: ${downloader}"

output="${output:-$("$downloader" --print-path windows11)}"
if [[ -e "$output" && $force -eq 0 ]]; then
  die "${output} exists (it may be the backing file of existing VM disks); use --force or another OUTPUT."
fi

case "${VM_IMAGE_ARCH:-$(uname -m)}" in
  arm64|aarch64) arch=arm64 ;;
  x86_64|amd64) arch=amd64 ;;
  *) die "unsupported architecture: ${VM_IMAGE_ARCH:-$(uname -m)}" ;;
esac

total_steps=5
step_number=0
step() {
  step_number=$((step_number + 1))
  printf '[%d/%d] %s\n' "$step_number" "$total_steps" "$*"
}
note() {
  printf '  %s\n' "$*" >&2
}

mkdir -p "${BASE_BUILD_TMPDIR:-/var/tmp}"
work="$(mktemp -d "${BASE_BUILD_TMPDIR:-/var/tmp}/build-windows-base.XXXXXX")"
pid_file="${work}/qemu.pid"
succeeded=0

ssh_options=(-o BatchMode=yes -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null
             -o LogLevel=ERROR -o ConnectTimeout=5)
[[ -z "${SSH_PRIVATE_KEY_FILE:-}" ]] || ssh_options+=(-i "$SSH_PRIVATE_KEY_FILE")
guest_ssh() { ssh -n "${ssh_options[@]}" -p "$ssh_port" guofang@127.0.0.1 "$@"; }

vm_running() {
  [[ -f "$pid_file" ]] && kill -0 "$(cat "$pid_file")" 2>/dev/null
}

cleanup() {
  if vm_running; then
    kill "$(cat "$pid_file")" 2>/dev/null || true
  fi
  if ((succeeded && !keep_work)); then
    rm -rf "$work"
  else
    note "Scratch directory kept: ${work}"
  fi
}
trap cleanup EXIT

# Poll "$@" every 10 s until it succeeds; give up after $1 seconds.
wait_for() {
  local limit="$1"
  shift
  local deadline=$((SECONDS + limit))
  until "$@"; do
    ((SECONDS < deadline)) || return 1
    sleep 10
  done
}

# Only the SSH rule is forwarded so that a running VM's RDP/DSH ports do not clash.
printf 'SSH  %s  22\n' "$ssh_port" >"${work}/forwards.conf"

step "Installing Windows 11 from the ISO (unattended; takes a while)"
display_option=(--vnc=127.0.0.1:90)
if ((window)); then
  [[ -n "${DISPLAY:-}${WAYLAND_DISPLAY:-}" || "$(uname -s)" == Darwin ]] \
    || die "--window needs a desktop session (DISPLAY or WAYLAND_DISPLAY is not set)."
  display_option=()
fi
WIN11_VM_FORWARDS="${work}/forwards.conf" "$launcher" \
  --daemon "${display_option[@]}" --install-from-iso --arch "$arch" --state-dir "$work"

step "Waiting for the guest to finish its first-logon setup (SSH answers)"
if ((window)); then
  note "Screen: QEMU window; console log: ${work}/console.log"
else
  note "Screen: VNC 127.0.0.1:5990; console log: ${work}/console.log"
fi
if ! wait_for "${BASE_BUILD_TIMEOUT:-5400}" guest_ssh 'echo ready' >/dev/null 2>&1; then
  die "SSH did not answer in time; inspect the guest over VNC."
fi

# The answer file the clones use on their first boot. generalize removes the
# machine identity but keeps the guofang account, so OOBE only needs the
# language, the EULA and the automatic logon; the specialize pass restores the
# settings Autounattend.xml applies (host name, time zone, Remote Desktop).
write_sysprep_answer() {
  cat >"$1" <<EOF
<?xml version="1.0" encoding="utf-8"?>
<unattend xmlns="urn:schemas-microsoft-com:unattend">
  <settings pass="specialize">
    <component name="Microsoft-Windows-Shell-Setup" processorArchitecture="${arch}"
      publicKeyToken="31bf3856ad364e35" language="neutral" versionScope="nonSxS">
      <ComputerName>WIN11</ComputerName>
      <TimeZone>China Standard Time</TimeZone>
    </component>
    <component name="Microsoft-Windows-TerminalServices-LocalSessionManager"
      processorArchitecture="${arch}" publicKeyToken="31bf3856ad364e35" language="neutral"
      versionScope="nonSxS">
      <fDenyTSConnections>false</fDenyTSConnections>
    </component>
  </settings>
  <settings pass="oobeSystem">
    <component name="Microsoft-Windows-International-Core" processorArchitecture="${arch}"
      publicKeyToken="31bf3856ad364e35" language="neutral" versionScope="nonSxS">
      <InputLocale>en-US</InputLocale>
      <SystemLocale>en-US</SystemLocale>
      <UILanguage>en-US</UILanguage>
      <UserLocale>en-US</UserLocale>
    </component>
    <component name="Microsoft-Windows-Shell-Setup" processorArchitecture="${arch}"
      publicKeyToken="31bf3856ad364e35" language="neutral" versionScope="nonSxS">
      <OOBE>
        <HideEULAPage>true</HideEULAPage>
        <HideOnlineAccountScreens>true</HideOnlineAccountScreens>
        <HideWirelessSetupInOOBE>true</HideWirelessSetupInOOBE>
        <SkipMachineOOBE>true</SkipMachineOOBE>
        <SkipUserOOBE>true</SkipUserOOBE>
        <ProtectYourPC>3</ProtectYourPC>
      </OOBE>
      <AutoLogon>
        <Enabled>true</Enabled>
        <Username>guofang</Username>
        <Password>
          <Value>wang</Value>
          <PlainText>true</PlainText>
        </Password>
      </AutoLogon>
    </component>
  </settings>
</unattend>
EOF
}

if ((sysprep)); then
  step "Generalizing the guest with sysprep"
  write_sysprep_answer "${work}/sysprep-unattend.xml"
  scp "${ssh_options[@]}" -P "$ssh_port" "${work}/sysprep-unattend.xml" \
    guofang@127.0.0.1:C:/Windows/Temp/sysprep-unattend.xml
  # The connection drops when the guest shuts down, which ssh reports as 255.
  status=0
  guest_ssh 'C:\Windows\System32\Sysprep\sysprep.exe /generalize /oobe /shutdown /quiet /unattend:C:\Windows\Temp\sysprep-unattend.xml' || status=$?
  if ((status != 0 && status != 255)); then
    guest_ssh 'type C:\Windows\System32\Sysprep\Panther\setuperr.log' >&2 || true
    die "sysprep failed (exit code ${status}); the log is above."
  fi
else
  step "Shutting the guest down (sysprep skipped)"
  guest_ssh 'shutdown /s /t 5' || true
fi

step "Waiting for the guest to power off"
if ! wait_for 900 bash -c '! kill -0 "$(cat "$1")" 2>/dev/null' _ "$pid_file"; then
  if ((sysprep)); then
    guest_ssh 'type C:\Windows\System32\Sysprep\Panther\setuperr.log' >&2 || true
  fi
  die "the guest did not power off; sysprep probably failed (log above, if SSH still answers)."
fi

step "Converting the disk to ${output}"
mkdir -p "$(dirname "$output")"
convert_options=(-f qcow2 -O qcow2)
((compress)) && convert_options+=(-c)
rm -f "${output}.part"
qemu-img convert -p "${convert_options[@]}" "${work}/disk.qcow2" "${output}.part"
[[ ! -e "$output" ]] || chmod u+w "$output"
mv -f "${output}.part" "$output"
chmod a-w "$output"
succeeded=1
printf 'Windows 11 base image is ready: %s\n' "$output"