#!/bin/bash

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(dirname "$SCRIPT_DIR")"

DISK_SIZE_GIB=96
MIN_FREE_GIB=64
RECREATE=0
GAP=0
ISO=""

ACTIVE_NBD=""
MOUNT_DIR=""
MOUNTED=0
TEMPORARY_DISK=""

usage() {
  cat <<USAGE
Usage: omarchy-iso-test-windows-disk [--recreate] [--gap] [release/omarchy.iso]

Creates or validates a synthetic Windows-style test disk with:
  - a 512MiB EFI partition containing EFI/Microsoft
  - an existing ext4 data partition with a preservation marker
  - more than 64GiB of contiguous unallocated space for Omarchy

This exercises Windows partition detection and preservation; it is not a
bootable Windows installation.

--gap builds the same fixture with a hole in the partition numbering: a third
partition is created and then deleted, leaving slot 3 free while slot 4 stays
in use. parted fills the lowest free slot, so the installer must create its
partitions as 3 and 5 and read those numbers back. The contiguous fixture
cannot catch that class of bug, because there "lowest free" and "highest + 1"
are the same number.

Options:
  --recreate  Replace the existing test fixture
  --gap       Build the numbering-hole variant (separate disk image)
  -h, --help  Show this help

Paths can be overridden with:
  OMARCHY_ISO_TEST_WINDOWS_DISK
  OMARCHY_ISO_TEST_WINDOWS_OVMF_VARS
USAGE
}

while (( $# > 0 )); do
  case "$1" in
    --recreate)
      RECREATE=1
      shift
      ;;
    --gap)
      GAP=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    -*)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 2
      ;;
    *)
      if [[ -n $ISO ]]; then
        echo "Only one ISO may be specified." >&2
        usage >&2
        exit 2
      fi
      ISO="$1"
      shift
      ;;
  esac
done

# Separate images per layout so the two fixtures can coexist and neither has
# to be rebuilt to run the other.
if (( GAP )); then
  TEST_DISK="${OMARCHY_ISO_TEST_WINDOWS_DISK:-/tmp/omarchy-iso-test-windows-gap.qcow2}"
  TEST_OVMF_VARS="${OMARCHY_ISO_TEST_WINDOWS_OVMF_VARS:-/tmp/OVMF_VARS.omarchy-iso-test-windows-gap.fd}"
  LAYOUT_LABEL="512MiB Windows ESP + data + deleted slot 3 + data + 64GiB+ free"
else
  TEST_DISK="${OMARCHY_ISO_TEST_WINDOWS_DISK:-/tmp/omarchy-iso-test-windows.qcow2}"
  TEST_OVMF_VARS="${OMARCHY_ISO_TEST_WINDOWS_OVMF_VARS:-/tmp/OVMF_VARS.omarchy-iso-test-windows.fd}"
  LAYOUT_LABEL="512MiB Windows ESP + existing data + 64GiB+ free"
fi

cleanup_mount() {
  if (( MOUNTED )); then
    sudo umount "$MOUNT_DIR" >/dev/null 2>&1 || true
    MOUNTED=0
  fi
  if [[ -n $MOUNT_DIR ]]; then
    rmdir "$MOUNT_DIR" >/dev/null 2>&1 || true
    MOUNT_DIR=""
  fi
}

disconnect_disk() {
  cleanup_mount
  if [[ -n $ACTIVE_NBD ]]; then
    sudo qemu-nbd --disconnect "$ACTIVE_NBD" >/dev/null 2>&1 || true
    ACTIVE_NBD=""
  fi
}

cleanup() {
  disconnect_disk
  [[ -z $TEMPORARY_DISK ]] || rm -f "$TEMPORARY_DISK"
}

trap cleanup EXIT
trap 'exit 130' INT
trap 'exit 143' TERM

find_free_nbd() {
  local device
  for device in /sys/class/block/nbd*; do
    [[ -e $device ]] || continue
    [[ ! -e $device/pid ]] || continue
    printf '/dev/%s\n' "${device##*/}"
    return 0
  done
  return 1
}

connect_disk() {
  local disk="$1" mode="${2:-read-write}"
  local -a options=(--format=qcow2)

  ACTIVE_NBD="$(find_free_nbd)" || {
    echo "No unused NBD device is available." >&2
    return 1
  }
  [[ $mode == read-only ]] && options+=(--read-only)

  if ! sudo qemu-nbd "${options[@]}" --connect="$ACTIVE_NBD" "$disk"; then
    ACTIVE_NBD=""
    return 1
  fi
  if ! sudo udevadm settle; then
    disconnect_disk
    return 1
  fi
}

mount_partition() {
  local partition="$1"
  MOUNT_DIR="$(mktemp -d)" || return 1
  if ! sudo mount -o "${2:-defaults}" "$partition" "$MOUNT_DIR"; then
    rmdir "$MOUNT_DIR" >/dev/null 2>&1 || true
    MOUNT_DIR=""
    return 1
  fi
  MOUNTED=1
}

largest_free_region() {
  sudo parted -ms "$1" unit B print free </dev/null 2>/dev/null |
    awk -F: '
      $NF ~ /free/ {
        size=$4
        gsub(/B/, "", size)
        if (size+0 > max+0) max=size
      }
      END { print max+0 }
    '
}

validate_test_disk() {
  local part_count pttype free_space_b min_free_b want_parts numbers
  local esp data valid=1

  [[ -f $TEST_DISK ]] || return 1
  qemu-img check -q "$TEST_DISK" >/dev/null 2>&1 || return 1
  connect_disk "$TEST_DISK" read-only || return 1

  esp="${ACTIVE_NBD}p1"
  data="${ACTIVE_NBD}p2"
  part_count=$(lsblk -nro TYPE "$ACTIVE_NBD" | awk '$1 == "part" { count++ } END { print count+0 }')
  pttype=$(lsblk -dnro PTTYPE "$ACTIVE_NBD" 2>/dev/null || true)
  free_space_b=$(largest_free_region "$ACTIVE_NBD")
  min_free_b=$((MIN_FREE_GIB * 1024 * 1024 * 1024))
  want_parts=2
  (( GAP )) && want_parts=3

  [[ $pttype == gpt ]] || valid=0
  (( part_count == want_parts )) || valid=0

  # The whole point of the gap fixture: slot 3 free, slot 4 in use.
  if (( GAP )); then
    numbers=$(sudo parted -ms "$ACTIVE_NBD" unit B print 2>/dev/null | tail -n +3 | cut -d: -f1 | sort | tr '\n' ' ')
    [[ $numbers == "1 2 4 " ]] || valid=0
  fi
  (( free_space_b >= min_free_b )) || valid=0
  [[ $(sudo blkid -s TYPE -o value "$esp" 2>/dev/null || true) == vfat ]] || valid=0
  [[ $(sudo blkid -s LABEL -o value "$esp" 2>/dev/null || true) == WINDOWS_ESP ]] || valid=0
  [[ $(sudo blkid -s TYPE -o value "$data" 2>/dev/null || true) == ext4 ]] || valid=0
  [[ $(sudo blkid -s LABEL -o value "$data" 2>/dev/null || true) == WINDOWS_DATA ]] || valid=0

  if (( valid )); then
    mount_partition "$esp" ro || valid=0
    if (( valid )); then
      sudo test -d "$MOUNT_DIR/EFI/Microsoft" || valid=0
    fi
    cleanup_mount
  fi

  if (( valid )); then
    mount_partition "$data" ro || valid=0
    if (( valid )); then
      sudo test -f "$MOUNT_DIR/OMARCHY-MUST-NOT-TOUCH.txt" || valid=0
    fi
    cleanup_mount
  fi

  disconnect_disk
  (( valid ))
}

create_test_disk() {
  local esp data

  mkdir -p "$(dirname "$TEST_DISK")"
  TEMPORARY_DISK="${TEST_DISK}.creating.$$"
  rm -f "$TEMPORARY_DISK"

  gum style --foreground 6 "Creating synthetic Windows test disk at $TEST_DISK"
  qemu-img create -f qcow2 "$TEMPORARY_DISK" "${DISK_SIZE_GIB}G"
  connect_disk "$TEMPORARY_DISK"

  sudo parted --script "$ACTIVE_NBD" mklabel gpt
  sudo parted --script "$ACTIVE_NBD" mkpart WINDOWS_ESP fat32 1MiB 513MiB
  sudo parted --script "$ACTIVE_NBD" set 1 esp on
  sudo parted --script "$ACTIVE_NBD" mkpart WINDOWS_DATA ext4 513MiB 20GiB

  if (( GAP )); then
    # Slot 3 is created and deleted, slot 4 stays. parted then fills 3 before
    # 5, so an installer that assumes "highest + 1" (which would be 5 and 6)
    # writes to device nodes that do not exist. The deleted region is
    # deliberately smaller than the trailing free space so the installer still
    # picks the tail as its largest free region.
    sudo parted --script "$ACTIVE_NBD" mkpart SCRATCH ext4 20GiB 28GiB
    sudo parted --script "$ACTIVE_NBD" mkpart WINDOWS_MORE ext4 28GiB 30GiB
    sudo parted --script "$ACTIVE_NBD" rm 3
  fi

  sudo partprobe "$ACTIVE_NBD"
  sudo udevadm settle

  esp="${ACTIVE_NBD}p1"
  data="${ACTIVE_NBD}p2"
  sudo mkfs.fat -F32 -n WINDOWS_ESP "$esp" >/dev/null
  sudo mkfs.ext4 -F -L WINDOWS_DATA "$data" >/dev/null 2>&1

  mount_partition "$esp"
  sudo mkdir -p "$MOUNT_DIR/EFI/Microsoft/Boot"
  printf '%s\n' 'Synthetic Windows ESP for Omarchy ISO testing.' |
    sudo tee "$MOUNT_DIR/EFI/Microsoft/README.txt" >/dev/null
  cleanup_mount

  mount_partition "$data"
  printf '%s\n' 'This existing partition must survive the Omarchy installation.' |
    sudo tee "$MOUNT_DIR/OMARCHY-MUST-NOT-TOUCH.txt" >/dev/null
  cleanup_mount

  sync
  disconnect_disk
  mv -f "$TEMPORARY_DISK" "$TEST_DISK"
  TEMPORARY_DISK=""
}

select_iso() {
  local -a isos=()

  if [[ -n $ISO ]]; then
    [[ -f $ISO ]] || {
      echo "ISO not found: $ISO" >&2
      return 1
    }
    realpath "$ISO"
    return 0
  fi

  mapfile -t isos < <(
    find "$REPO_ROOT/release" -maxdepth 1 -type f -name '*.iso' -printf '%T@ %p\n' 2>/dev/null |
      sort -nr |
      cut -d' ' -f2-
  )
  (( ${#isos[@]} > 0 )) || {
    echo "No ISO found under $REPO_ROOT/release." >&2
    return 1
  }

  printf '%s\n' "${isos[@]}" | gum choose --header "Select ISO to install"
}

omarchy-pkg-add qemu-full edk2-ovmf parted dosfstools e2fsprogs
sudo modprobe nbd max_part=16

if [[ -f $TEST_DISK ]] && (( ! RECREATE )); then
  gum style --foreground 6 "Verifying existing test disk at $TEST_DISK"
  if validate_test_disk; then
    gum style --foreground 2 "✓ Synthetic Windows test disk is ready"
  elif gum confirm --default=false "The existing test disk is not a pristine fixture. Recreate it?"; then
    create_test_disk
  else
    gum style "Left the existing disk untouched."
    exit 0
  fi
else
  create_test_disk
fi

# The synthetic fixture has no real Windows bootloader, so reset OVMF state for
# every run. The ISO still sees EFI/Microsoft and must leave that ESP untouched.
mkdir -p "$(dirname "$TEST_OVMF_VARS")"
cp /usr/share/edk2/x64/OVMF_VARS.4m.fd "$TEST_OVMF_VARS"

gum style --border rounded --padding "0 1" --border-foreground 6 \
  "Synthetic Windows disk ready" \
  "Disk: $TEST_DISK" \
  "Layout: $LAYOUT_LABEL"

if gum confirm "Start installation on this disk?"; then
  ISO="$(select_iso)"
  trap - EXIT INT TERM
  exec "$SCRIPT_DIR/omarchy-iso-boot" \
    --reuse \
    --disk "$TEST_DISK" \
    --ovmf-vars "$TEST_OVMF_VARS" \
    "$ISO" \
    -- -boot once=d,menu=on
fi
