#!/usr/bin/env bash
# Headless Windows Server 2025 Core test VMs (QEMU/KVM).
#
#   vm.sh create                 download the evaluation ISO, install base.qcow2
#   vm.sh variants               en.qcow2 / zh.qcow2 overlays on top of base.qcow2
#                                (zh: zh-CN system locale, ANSI code page 936)
#   vm.sh start|stop|status|wait VARIANT
#   vm.sh ssh VARIANT [USER] [command...]   USER: Administrator (default), tester, 测试
#   vm.sh put VARIANT SRC... DEST           scp as Administrator
#   vm.sh ps VARIANT [USER] FILE.ps1 [args...]   run a script (uploaded with a UTF-8 BOM)
#   vm.sh shot VARIANT [out.png]            console screenshot
#   vm.sh type VARIANT TEXT                 type into the console (\n = Enter)
#   vm.sh login VARIANT                     unlock the console as Administrator
#
# State (ISO, disks, keys, passwords) lives in $CST_WINVM_DIR, default ~/winvm.
# Keys and passwords are created there with mode 0600 and never printed.
set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
D="${CST_WINVM_DIR:-$HOME/winvm}"
ISO_URL="https://aka.ms/WinServ2025iso-enus"
IMAGE_NAME="Windows Server 2025 SERVERSTANDARDCORE"
USERS=(tester 测试)

if ! command -v qemu-system-x86_64 >/dev/null || ! command -v sshpass >/dev/null ||
  ! command -v socat >/dev/null || ! command -v xorriso >/dev/null; then
  if command -v nix-shell >/dev/null && [[ -z "${CST_WINVM_NIX:-}" ]]; then
    export CST_WINVM_NIX=1
    exec nix-shell -p qemu sshpass socat xorriso imagemagick python3 \
      --run "$(printf '%q ' "$0" "$@")"
  fi
  echo "need qemu-system-x86_64, sshpass, socat, xorriso" >&2
  exit 1
fi

die() { echo "vm.sh: $*" >&2; exit 1; }

port_of() {
  case "$1" in
  base) echo 2221 ;;
  en) echo 2222 ;;
  zh) echo 2223 ;;
  *) die "unknown variant '$1' (base, en, zh)" ;;
  esac
}

pw_file() {
  case "$1" in
  Administrator) echo "$D/admin-password" ;;
  tester) echo "$D/tester-password" ;;
  测试) echo "$D/ceshi-password" ;;
  *) die "unknown user '$1'" ;;
  esac
}

new_password() { # prefix -> file, mixed classes for Windows complexity rules
  local f="$2"
  [[ -s "$f" ]] && return 0
  (umask 077; printf '%s%s9a' "$1" "$(head -c 256 /dev/urandom | tr -dc 'A-Za-z0-9' | head -c 16)" >"$f")
}

running() {
  local pid
  pid="$(cat "$D/$1.pid" 2>/dev/null)" || return 1
  [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null
}

ssh_opts() {
  echo -p "$(port_of "$1")" -o UserKnownHostsFile="$D/known_hosts" \
    -o StrictHostKeyChecking=accept-new -o LogLevel=ERROR -o ConnectTimeout=10
}

vm_ssh() { # variant user command...
  local v="$1" u="$2"
  shift 2
  # shellcheck disable=SC2046
  if [[ "$u" == Administrator ]]; then
    ssh -i "$D/vm_key" -o BatchMode=yes $(ssh_opts "$v") "Administrator@127.0.0.1" "$@"
  else
    SSHPASS="$(cat "$(pw_file "$u")")" sshpass -e ssh -o PubkeyAuthentication=no \
      -o PreferredAuthentications=password $(ssh_opts "$v") "$u@127.0.0.1" "$@"
  fi
}

vm_put() { # variant src... dest
  local v="$1"
  shift
  local dest="${!#}"
  local srcs=("${@:1:$#-1}")
  # shellcheck disable=SC2046
  scp -q -r -i "$D/vm_key" -o BatchMode=yes $(ssh_opts "$v" | sed 's/^-p/-P/') \
    "${srcs[@]}" "Administrator@127.0.0.1:$dest"
}

qemu_start() { # variant disk [extra args...]
  local v="$1" disk="$2"
  shift 2
  running "$v" && return 0
  [[ -f "$disk" ]] || die "missing $disk"
  qemu-system-x86_64 \
    -name "cst-win-$v" -machine q35,accel=kvm \
    -cpu host,hv_relaxed,hv_spinlocks=0x1fff,hv_vapic,hv_time \
    -smp 4 -m 6G -rtc base=localtime \
    -device ahci,id=ahci \
    -drive file="$disk",if=none,id=d0,cache=writeback,discard=unmap \
    -device ide-hd,drive=d0,bus=ahci.0 \
    -netdev user,id=n0,hostfwd=tcp:127.0.0.1:"$(port_of "$v")"-:22 \
    -device e1000e,netdev=n0 \
    -device qemu-xhci,id=xhci -device usb-tablet,bus=xhci.0 \
    -display none -qmp unix:"$D/$v.qmp",server,nowait \
    -pidfile "$D/$v.pid" -daemonize "$@"
}

disk_of() {
  case "$1" in
  base) echo "$D/base.qcow2" ;;
  *) echo "$D/$1.qcow2" ;;
  esac
}

vm_wait() { # variant [timeout_s]
  local v="$1" t="${2:-900}" i
  for ((i = 0; i < t; i += 5)); do
    if vm_ssh "$v" Administrator 'Test-Path C:\setup.done' 2>/dev/null | grep -q True; then
      return 0
    fi
    running "$v" || die "$v is not running"
    sleep 5
  done
  die "$v: no SSH after ${t}s"
}

vm_stop() {
  local v="$1" i
  running "$v" || return 0
  vm_ssh "$v" Administrator 'Stop-Computer -Force' >/dev/null 2>&1 || true
  for ((i = 0; i < 120; i++)); do
    running "$v" || { rm -f "$D/$v.pid" "$D/$v.qmp"; return 0; }
    sleep 1
  done
  printf '%s\n' '{"execute":"qmp_capabilities"}' '{"execute":"quit"}' |
    socat - UNIX-CONNECT:"$D/$v.qmp" >/dev/null || true
  rm -f "$D/$v.pid" "$D/$v.qmp"
}

with_bom() { # file -> temp copy with a UTF-8 BOM (Windows PowerShell 5.1 reads ANSI otherwise)
  local out
  out="$(mktemp --suffix=.ps1)"
  printf '\xef\xbb\xbf' >"$out"
  sed '1s/^\xef\xbb\xbf//' "$1" >>"$out"
  echo "$out"
}

vm_ps() { # variant user file args...
  local v="$1" u="$2" f="$3"
  shift 3
  local tmp name
  tmp="$(with_bom "$f")"
  name="cst-$(basename "$f" .ps1)-$$.ps1"
  vm_put "$v" "$tmp" "C:/Windows/Temp/$name"
  rm -f "$tmp"
  vm_ssh "$v" Administrator "icacls C:\\Windows\\Temp\\$name /grant 'Users:RX' | Out-Null" >/dev/null
  local rc=0
  vm_ssh "$v" "$u" "powershell -NoProfile -ExecutionPolicy Bypass -File C:\\Windows\\Temp\\$name $*" || rc=$?
  vm_ssh "$v" Administrator "Remove-Item -Force C:\\Windows\\Temp\\$name" >/dev/null 2>&1 || true
  return $rc
}

provision() { # variant: users, sshd, console; idempotent
  local v="$1" u tmp
  tmp="$(mktemp)"
  chmod 600 "$tmp"
  for u in "${USERS[@]}"; do
    new_password "Tst!" "$(pw_file "$u")"
    printf '%s\t%s\n' "$u" "$(cat "$(pw_file "$u")")" >>"$tmp"
  done
  vm_put "$v" "$tmp" "C:/Windows/Temp/cst-users.txt"
  rm -f "$tmp"
  vm_ps "$v" Administrator "$HERE/provision.ps1"
}

cmd_create() {
  mkdir -p "$D"
  chmod 700 "$D"
  [[ -f "$D/base.qcow2" ]] && die "$D/base.qcow2 exists"
  if [[ ! -s "$D/server2025.iso" ]]; then
    echo "==> downloading the Windows Server 2025 evaluation ISO (~8 GB)"
    curl -fL -C - -o "$D/server2025.iso" "$ISO_URL"
  fi
  [[ -f "$D/vm_key" ]] || ssh-keygen -q -t ed25519 -N "" -C cst-winvm -f "$D/vm_key"
  new_password "Cst!" "$D/admin-password"

  local u
  u="$(mktemp -d)"
  chmod 700 "$u"
  cp "$HERE/setup.ps1" "$D/vm_key.pub" "$u/"
  (umask 077
    sed -e "s|@IMAGE_NAME@|$IMAGE_NAME|" -e "s|@ADMIN_PASSWORD@|$(cat "$D/admin-password")|g" \
      "$HERE/autounattend.xml.in" >"$u/autounattend.xml"
    xorriso -as mkisofs -quiet -J -r -V UNATTEND -o "$D/unattend.iso" "$u/" 2>/dev/null)
  rm -rf "$u"

  qemu-img create -q -f qcow2 "$D/base.qcow2" 60G
  echo "==> unattended install (20-30 minutes)"
  qemu_start base "$D/base.qcow2" \
    -drive file="$D/server2025.iso",media=cdrom,if=none,id=cd0,readonly=on \
    -device ide-cd,drive=cd0,bus=ahci.1 \
    -drive file="$D/unattend.iso",media=cdrom,if=none,id=cd1,readonly=on \
    -device ide-cd,drive=cd1,bus=ahci.2 \
    -boot order=c,once=d
  vm_wait base 3600
  provision base
  vm_stop base
  echo "OK: $D/base.qcow2; next: vm.sh variants"
}

cmd_variants() {
  [[ -f "$D/base.qcow2" ]] || die "no base.qcow2; run vm.sh create"
  running base && die "stop the base VM first"
  local v
  for v in en zh; do
    running "$v" && vm_stop "$v"
    rm -f "$D/$v.qcow2"
    qemu-img create -q -f qcow2 -b "$D/base.qcow2" -F qcow2 "$D/$v.qcow2"
  done
  qemu_start zh "$D/zh.qcow2"
  vm_wait zh
  vm_ssh zh Administrator 'Set-WinSystemLocale -SystemLocale zh-CN; Set-Culture zh-CN; Set-WinHomeLocation -GeoId 45; Restart-Computer -Force' || true
  sleep 20
  vm_wait zh
  echo "zh ANSI code page: $(vm_ssh zh Administrator '[Text.Encoding]::Default.CodePage')"
  vm_stop zh
  echo "OK: en.qcow2 zh.qcow2"
}

cmd="${1:-}"
shift || true
case "$cmd" in
create) cmd_create ;;
variants) cmd_variants ;;
provision) provision "${1:?variant}" ;;
start)
  qemu_start "${1:?variant}" "$(disk_of "$1")"
  vm_wait "$1"
  ;;
stop) vm_stop "${1:?variant}" ;;
status) running "${1:?variant}" && echo running || echo stopped ;;
wait) vm_wait "${1:?variant}" "${2:-900}" ;;
ssh)
  v="${1:?variant}"
  shift
  u=Administrator
  if [[ $# -gt 0 ]] && [[ "$1" == Administrator || "$1" == tester || "$1" == 测试 ]]; then
    u="$1"
    shift
  fi
  vm_ssh "$v" "$u" "$@"
  ;;
put) vm_put "$@" ;;
ps)
  v="${1:?variant}"
  shift
  u=Administrator
  if [[ "$1" == Administrator || "$1" == tester || "$1" == 测试 ]]; then
    u="$1"
    shift
  fi
  vm_ps "$v" "$u" "$@"
  ;;
shot)
  v="${1:?variant}"
  out="${2:-$D/$v.png}"
  printf '%s\n' '{"execute":"qmp_capabilities"}' \
    '{"execute":"screendump","arguments":{"filename":"'"$D/$v.ppm"'"}}' |
    socat - UNIX-CONNECT:"$D/$v.qmp" >/dev/null
  sleep 1
  magick "$D/$v.ppm" -resize '1024x1024>' "$out"
  rm -f "$D/$v.ppm"
  echo "$out"
  ;;
type) python3 "$HERE/type.py" "$D/${1:?variant}.qmp" "${2:?text}" ;;
login)
  python3 "$HERE/type.py" "$D/${1:?variant}.qmp" --cad
  sleep 3
  python3 "$HERE/type.py" "$D/$1.qmp" "@$D/admin-password"
  ;;
*) sed -n '2,15p' "$0" | sed 's/^# \{0,1\}//'; exit 2 ;;
esac