#!/usr/bin/env bash
# Build a bootable Windows 11 installation disk image from the Windows ISO.
#
# Attached to the VM as an NVMe disk instead of a USB CD-ROM, the installation
# files are read about 15 times faster (QEMU's emulated USB storage reads the ISO
# at a few MB/s). The image holds one FAT32 partition, which the UEFI firmware
# can boot. FAT32 files are limited to 4 GB, so sources/install.wim is split into
# install.swm, install2.swm and so on, which Windows Setup reads as one image.
#
# Runs on macOS and Linux (Debian) and needs no root for the image itself: the
# image is filled with mtools. The ISO is a hybrid of ISO 9660 and UDF whose
# ISO 9660 side only has a README, so it must be read as UDF:
#   macOS            hdiutil attach
#   Linux            udisksctl (udisks2)
#
# Dependencies
#   macOS (Homebrew)  wimlib dosfstools gptfdisk mtools
#   Debian            wimtools dosfstools gdisk mtools udisks2
#
# Usage: build-windows-install-disk.sh [--force] ISO [OUTPUT.img]
# OUTPUT defaults to the ISO path with ".install.img" in place of ".iso".
set -euo pipefail

usage() {
  printf 'Usage: %s [--force] ISO [OUTPUT.img]\n' "$(basename "$0")"
}

force=0
args=()
for arg in "$@"; do
  case "$arg" in
    -h|--help) usage; exit 0 ;;
    --force) force=1 ;;
    *) args+=("$arg") ;;
  esac
done
if ((${#args[@]} < 1 || ${#args[@]} > 2)); then
  usage >&2
  exit 2
fi
iso="${args[0]}"
output="${args[1]:-${iso%.*}.install.img}"

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

# The build has four steps; step prints a numbered heading for each of them.
total_steps=4
step_number=0
step() {
  step_number=$((step_number + 1))
  printf '[%d/%d] %s\n' "$step_number" "$total_steps" "$*"
}

# Progress details under a step; they go to stderr to keep stdout for the
# headings and the result line.
note() {
  printf '  %s\n' "$*" >&2
}

# show logs a command line before it runs; run logs it and then runs it. Use show
# alone where the output of the command is captured.
show() {
  local line
  printf -v line '%q ' "$@"
  note "\$ ${line% }"
}
run() {
  show "$@"
  "$@"
}

os="$(uname -s)"
case "$os" in
  Darwin|Linux) ;;
  *) die "unsupported system: $os" ;;
esac

# mkfs.fat and sgdisk live in sbin, which is usually not in a user's PATH (and
# Homebrew has its own sbin).
brew_prefix=""
if [[ "$os" == Darwin ]] && command -v brew >/dev/null 2>&1; then
  brew_prefix="$(brew --prefix 2>/dev/null || true)"
fi
for dir in ${brew_prefix:+"$brew_prefix/sbin"} /opt/homebrew/sbin /usr/local/sbin /usr/sbin /sbin; do
  [[ -d "$dir" ]] && PATH="$PATH:$dir"
done

tools=(wimlib-imagex mkfs.fat sgdisk mcopy mmd)
[[ "$os" == Darwin ]] && tools+=(hdiutil)
missing=()
for tool in "${tools[@]}"; do
  command -v "$tool" >/dev/null 2>&1 || missing+=("$tool")
done
if ((${#missing[@]} > 0)); then
  if [[ "$os" == Darwin ]]; then
    hint="brew install wimlib dosfstools gptfdisk mtools"
  else
    hint="apt install wimtools dosfstools gdisk mtools"
  fi
  die "missing tools: ${missing[*]} ($hint)"
fi

[[ -f "$iso" ]] || die "ISO was not found: $iso"
if [[ -e "$output" && $force -eq 0 ]]; then
  die "$output already exists (use --force to replace it)"
fi

# stat prints a file size with -f %z on macOS (BSD) and -c %s on Linux (GNU).
if [[ "$os" == Darwin ]]; then size_args=(-f %z); else size_args=(-c %s); fi

work_dir="$(mktemp -d "${TMPDIR:-/tmp}/win11-install-disk.XXXXXX")"
mount_dir="$work_dir/iso"
mkdir "$mount_dir"
note "work directory: $work_dir"
mount_method=""
loop_device=""
# Set while the output image is being written, so that a failed or cancelled
# run does not leave a half-built image behind.
partial_output=""
cleanup() {
  [[ -z "$partial_output" ]] || rm -f "$partial_output"
  case "$mount_method" in
    hdiutil) hdiutil detach -quiet "$mount_dir" >/dev/null 2>&1 || true ;;
    udisks)
      udisksctl unmount -b "$loop_device" >/dev/null 2>&1 || true
      udisksctl loop-delete -b "$loop_device" >/dev/null 2>&1 || true ;;
  esac
  rm -rf "$work_dir"
}
trap cleanup EXIT
# Turn the signals of a cancelled run into a normal exit, so that the EXIT trap
# unmounts the ISO and removes the work directory.
trap 'exit 130' INT
trap 'exit 143' TERM
trap 'exit 129' HUP

mount_with_udisks() {
  if ! command -v udisksctl >/dev/null 2>&1; then
    note "udisksctl is not installed"
    return 1
  fi
  local setup
  show udisksctl loop-setup -r -f "$iso"
  if ! setup="$(udisksctl loop-setup -r -f "$iso" 2>&1)"; then
    note "udisksctl loop-setup failed: $setup"
    return 1
  fi
  loop_device="$(sed -n 's/^Mapped file .* as \(\/dev\/loop[0-9]*\)\.$/\1/p' <<<"$setup")"
  if [[ -z "$loop_device" ]]; then
    note "could not find the loop device in: $setup"
    return 1
  fi
  note "loop device: $loop_device"
  # Set at once so that a cancel from here on still releases the loop device.
  mount_method=udisks
  # A desktop session may have mounted the loop device already, which makes
  # udisksctl mount fail with AlreadyMounted; use that mount point then.
  local mounted point
  show lsblk -no MOUNTPOINT "$loop_device"
  point="$(lsblk -no MOUNTPOINT "$loop_device" 2>/dev/null | head -n 1)"
  if [[ -z "$point" ]]; then
    show udisksctl mount -b "$loop_device" -t udf -o ro
    if mounted="$(udisksctl mount -b "$loop_device" -t udf -o ro 2>&1)"; then
      # Some udisks versions end the message with a period, others do not.
      point="$(sed -n 's/^Mounted .* at \(.*\)$/\1/p' <<<"$mounted")"
      [[ -d "$point" ]] || point="${point%.}"
      note "udisksctl mount said: $mounted"
    else
      note "udisksctl mount failed: $mounted"
    fi
  else
    note "$loop_device was already mounted"
  fi
  if [[ ! -d "$point" ]]; then
    note "no usable mount point for $loop_device (parsed mount point: '${point}')"
    run udisksctl loop-delete -b "$loop_device" >/dev/null 2>&1 || true
    mount_method=""
    return 1
  fi
  mount_dir="$point"
  note "mounted at $mount_dir"
}

step "Reading the ISO $iso (it is UDF, so it has to be mounted)"
note "ISO size: $(stat -L "${size_args[@]}" "$iso") bytes"
if [[ "$os" == Darwin ]]; then
  note "mounting with hdiutil at $mount_dir"
  run hdiutil attach -readonly -nobrowse -mountpoint "$mount_dir" "$iso" >/dev/null
  mount_method=hdiutil
else
  note "mounting with udisksctl"
  mount_with_udisks || true
fi
note "read method: ${mount_method:-none}"
[[ -n "$mount_method" ]] || die "cannot mount the ISO as UDF; install udisks2 and run from a desktop or login session"

install_wim="$mount_dir/sources/install.wim"
[[ -f "$install_wim" ]] || die "sources/install.wim was not found in the ISO"
[[ -d "$mount_dir/efi/boot" ]] || die "efi/boot was not found in the ISO; is it a UEFI Windows ISO?"
note "install.wim: $(stat "${size_args[@]}" "$install_wim") bytes"

# Parts must stay below the 4 GiB FAT32 file limit; wimlib takes MiB.
step "Splitting install.wim into install.swm parts (FAT32 files must stay below 4 GiB)"
split_dir="$work_dir/split"
mkdir "$split_dir"
run wimlib-imagex split "$install_wim" "$split_dir/install.swm" 3800 >/dev/null
for part in "$split_dir"/install*.swm; do
  note "$(basename "$part"): $(stat "${size_args[@]}" "$part") bytes"
done

# Size the image: everything but install.wim, the split parts, and 5% plus 64 MiB
# of room for FAT32 metadata and partition alignment.
content_bytes=$(
  { find "$mount_dir" -type f ! -path "$install_wim" -exec stat "${size_args[@]}" {} +
    find "$split_dir" -type f -exec stat "${size_args[@]}" {} +; } |
    awk '{ total += $1 } END { printf "%d", total }'
)
mib=$((1024 * 1024))
image_mib=$(( (content_bytes + content_bytes / 20) / mib + 64 ))
step "Creating a $image_mib MiB image with a GPT and one FAT32 partition"
note "content to copy: $content_bytes bytes (without install.wim)"

image="$output"
run mkdir -p "$(dirname -- "$output")"
# --force replaces an existing image; truncate alone would keep its old content.
run rm -f "$image"
partial_output="$image"
run truncate -s "${image_mib}m" "$image"
run sgdisk -o -n 1:2048:0 -t 1:0700 -c 1:WINPE "$image" >/dev/null
show sgdisk -i 1 "$image"
sectors=$(sgdisk -i 1 "$image" | awk '/^Partition size:/ { print $3 }')
[[ -n "$sectors" ]] || die "could not read the partition size"
run mkfs.fat -F 32 -n WINPE --offset 2048 "$image" $((sectors / 2)) >/dev/null

# mtools addresses the partition by its byte offset inside the image.
mtools_image="$image@@$((2048 * 512))"
note "partition: $sectors sectors, formatted as FAT32 (label WINPE)"
export MTOOLS_SKIP_CHECK=1

copy_dir() { # copy_dir SOURCE_DIR TARGET_DIR_IN_IMAGE
  local source="$1" target="$2" files=() entry
  while IFS= read -r -d '' entry; do
    files+=("$entry")
  done < <(find "$source" -maxdepth 1 -type f ! -path "$install_wim" -print0)
  if ((${#files[@]} > 0)); then
    note "\$ mcopy -Q -m -i $mtools_image <${#files[@]} files from ${source#"$mount_dir"} > ::$target/"
    mcopy -Q -m -i "$mtools_image" "${files[@]}" "::$target/"
  fi
}

step "Copying the installation files into the image"
copy_dir "$mount_dir" ""
dir_count=0
while IFS= read -r -d '' dir; do
  relative="${dir#"$mount_dir"/}"
  dir_count=$((dir_count + 1))
  run mmd -i "$mtools_image" "::$relative"
  copy_dir "$dir" "$relative"
done < <(find "$mount_dir" -mindepth 1 -type d -print0 | sort -z)
note "copied $dir_count directories"
note "copying the install*.swm parts to sources/"
run mcopy -Q -m -i "$mtools_image" "$split_dir"/install*.swm "::sources/"

partial_output=""
printf 'Built Windows 11 installation disk: %s\n' "$output"