#!/usr/bin/env bash
set -euo pipefail

images_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
base_qcow2_dir="${images_dir}/base-qcow2"
install_iso_dir="${images_dir}/install-iso"
tools_iso_dir="${images_dir}/tools-iso"

normalize_arch() {
  case "$1" in
    arm64|aarch64) printf 'arm64\n' ;;
    x86_64|amd64) printf 'amd64\n' ;;
    *) printf 'Error: unsupported architecture: %s\n' "$1" >&2; exit 1 ;;
  esac
}
image_arch="$(normalize_arch "${VM_IMAGE_ARCH:-$(uname -m)}")"

if [[ "$image_arch" == arm64 ]]; then
  debian_image="${base_qcow2_dir}/debian-13-generic-arm64-20260810-2566.qcow2"
  alma_image="${base_qcow2_dir}/AlmaLinux-10-GenericCloud-10.2-20260526.0.aarch64.qcow2"
  fedora_image="${base_qcow2_dir}/Fedora-Cloud-Base-Generic-44-1.7.aarch64.qcow2"
  debian_url_default="https://cloud.debian.org/images/cloud/trixie/20260810-2566/debian-13-generic-arm64-20260810-2566.qcow2"
  alma_url_default="https://repo.almalinux.org/almalinux/10/cloud/aarch64/images/AlmaLinux-10-GenericCloud-10.2-20260526.0.aarch64.qcow2"
  fedora_url_default="https://mirrors.tuna.tsinghua.edu.cn/fedora/releases/44/Cloud/aarch64/images/Fedora-Cloud-Base-Generic-44-1.7.aarch64.qcow2"
  # TODO: these were pinned for the amd64/x86_64 builds; verification is
  # skipped (empty default) for the arm64/aarch64 builds until their
  # checksums are fetched and pinned via DEBIAN_IMAGE_SHA512 /
  # ALMA_IMAGE_SHA256 / FEDORA_IMAGE_SHA256.
  debian_sha512_default=""
  alma_sha256_default=""
  fedora_sha256_default=""
else
  debian_image="${base_qcow2_dir}/debian-13-generic-amd64-20260810-2566.qcow2"
  alma_image="${base_qcow2_dir}/AlmaLinux-10-GenericCloud-10.2-20260526.0.x86_64.qcow2"
  fedora_image="${base_qcow2_dir}/Fedora-Cloud-Base-Generic-44-1.7.x86_64.qcow2"
  debian_url_default="https://cloud.debian.org/images/cloud/trixie/20260810-2566/debian-13-generic-amd64-20260810-2566.qcow2"
  alma_url_default="https://repo.almalinux.org/almalinux/10/cloud/x86_64/images/AlmaLinux-10-GenericCloud-10.2-20260526.0.x86_64.qcow2"
  fedora_url_default="https://mirrors.tuna.tsinghua.edu.cn/fedora/releases/44/Cloud/x86_64/images/Fedora-Cloud-Base-Generic-44-1.7.x86_64.qcow2"
  debian_sha512_default="f6978100d8031c266d55d7815ceea7fcdeacf28e1e5834fdb9c94ac96880a054a6e6f8681c2d3b0584e0057eaf3ef7353856b85212d04134744faa9b3bb1f24f"
  alma_sha256_default="47f2218668dd4776be140dd92fa3bea700be1766e2c7d88bdfd6a4b50f477b4d"
  fedora_sha256_default="28680fe5b371a5a82ebf43a31926e086a168e59949d03969c5093e7071f90b7f"
fi

virtio_image="${tools_iso_dir}/virtio-win.iso"

# Find a Windows 11 x64 file in a directory by matching "Win", "11", "x64" and
# the extension in that order (case-insensitive). When several match, the
# highest version-sorted name wins.
find_windows_image() {
  local dir="$1"
  local extension="$2"
  local match

  match="$(
    shopt -s nullglob nocaseglob
    files=("${dir}"/*Win*11*x64*."${extension}")
    ((${#files[@]} == 0)) || printf '%s\n' "${files[@]}" | sort -V | tail -n 1
  )"
  [[ -n "$match" ]] && printf '%s\n' "$match"
}

windows_iso_path() {
  if ! find_windows_image "$install_iso_dir" iso; then
    printf 'Error: no Windows 11 x64 ISO (*Win*11*x64*.iso) found in %s\n' "$install_iso_dir" >&2
    return 1
  fi
}

# The QCOW2 is named after the installation ISO. If the ISO has been removed,
# fall back to an existing matching QCOW2.
windows_qcow2_path() {
  local iso
  if iso="$(find_windows_image "$install_iso_dir" iso)"; then
    iso="$(basename -- "$iso")"
    printf '%s/%s.qcow2\n' "$base_qcow2_dir" "${iso%.*}"
  elif ! find_windows_image "$base_qcow2_dir" qcow2; then
    printf 'Error: no Windows 11 x64 ISO (*Win*11*x64*.iso) found in %s\n' "$install_iso_dir" >&2
    return 1
  fi
}

debian_url="${DEBIAN_IMAGE_URL:-$debian_url_default}"
alma_url="${ALMA_IMAGE_URL:-$alma_url_default}"
fedora_url="${FEDORA_IMAGE_URL:-$fedora_url_default}"
virtio_url="${VIRTIO_WIN_ISO_URL:-https://fedorapeople.org/groups/virt/virtio-win/direct-downloads/stable-virtio/virtio-win.iso}"
debian_sha512="${DEBIAN_IMAGE_SHA512:-$debian_sha512_default}"
alma_sha256="${ALMA_IMAGE_SHA256:-$alma_sha256_default}"
fedora_sha256="${FEDORA_IMAGE_SHA256:-$fedora_sha256_default}"
force_download=false

usage() {
  cat <<EOF
Usage: $(basename "$0") [--force] [all|debian13|almalinux10|fedora44|virtio-win] ...
       $(basename "$0") --print-path TARGET [--iso]

Download one or more shared VM images. With no arguments, all targets are
processed. The Debian/AlmaLinux/Fedora image architecture (arm64 or amd64)
is auto-detected from the host; override it with VM_IMAGE_ARCH.

--force removes existing files and partial downloads for the selected targets,
then downloads them again from the beginning.
--print-path prints only the absolute image path without downloading it.
For windows11 it prints the QCOW2 path by default; add --iso for the ISO path.
The Windows ISO is found automatically in install-iso/ by matching
*Win*11*x64*.iso; the QCOW2 path is named after that ISO.
EOF
}

print_image_path() {
  local target="$1"
  local media="${2:-qcow2}"

  case "$target" in
    debian13) printf '%s\n' "$debian_image" ;;
    almalinux10) printf '%s\n' "$alma_image" ;;
    fedora44) printf '%s\n' "$fedora_image" ;;
    virtio-win) printf '%s\n' "$virtio_image" ;;
    windows11)
      if [[ "$media" == iso ]]; then
        windows_iso_path
      else
        windows_qcow2_path
      fi
      ;;
    *)
      printf 'Error: path is unavailable for target: %s\n' "$1" >&2
      usage >&2
      return 2
      ;;
  esac
}

if [[ "${1:-}" == --print-path ]]; then
  if (($# < 2 || $# > 3)); then
    printf 'Error: --print-path requires one target and optionally --iso.\n' >&2
    usage >&2
    exit 2
  fi
  if (($# == 3)) && [[ "$2" != windows11 || "$3" != --iso ]]; then
    printf 'Error: --iso is only valid after the windows11 target.\n' >&2
    usage >&2
    exit 2
  fi
  media_option="${3:-}"
  print_image_path "$2" "${media_option#--}"
  exit
fi

targets=()
for argument in "$@"; do
  if [[ "$argument" == --force ]]; then
    force_download=true
  else
    targets+=("$argument")
  fi
done
if ((${#targets[@]} == 0)); then
  targets=(all)
fi

for target in "${targets[@]}"; do
  case "$target" in
    all|debian13|almalinux10|fedora44|virtio-win) ;;
    -h|--help) usage; exit 0 ;;
    *) printf 'Error: unknown download target: %s\n' "$target" >&2; usage >&2; exit 2 ;;
  esac
done

wants_target() {
  local requested="$1"
  local target
  for target in "${targets[@]}"; do
    [[ "$target" == all || "$target" == "$requested" ]] && return 0
  done
  return 1
}

mkdir -p "$base_qcow2_dir" "$install_iso_dir" "$tools_iso_dir"

verify_file() {
  local label="$1"
  local file="$2"
  local algorithm="$3"
  local expected="$4"
  local command_name="${algorithm}sum"
  local actual

  if ! command -v "$command_name" >/dev/null 2>&1; then
    printf 'Error: %s is required to verify %s.\n' "$command_name" "$label" >&2
    return 1
  fi

  actual="$("$command_name" "$file")"
  actual="${actual%% *}"
  if [[ "$actual" != "$expected" ]]; then
    printf 'Error: %s checksum mismatch for %s.\nFile:     %s\nExpected: %s\nActual:   %s\n' \
      "${algorithm^^}" "$label" "$file" "$expected" "$actual" >&2
    return 1
  fi

  printf 'Verified %s checksum for %s.\n' "${algorithm^^}" "$label"
}

download_file() {
  local label="$1"
  local url="$2"
  local destination="$3"
  local checksum_algorithm="${4:-}"
  local expected_checksum="${5:-}"
  local partial="${destination}.part"
  local partial_url="${partial}.url"
  local downloader

  if [[ "$force_download" == true ]]; then
    printf 'Forcing a fresh download of %s.\n' "$label"
    rm -f -- "$destination" "$partial" "$partial_url"
  fi

  if [[ -s "$destination" ]]; then
    if [[ -n "$expected_checksum" ]]; then
      verify_file "$label" "$destination" "$checksum_algorithm" "$expected_checksum" || return 1
    fi
    printf '%s already exists; skipping:\n%s\n' "$label" "$destination"
    return
  fi

  printf 'Downloading %s from:\n%s\n' "$label" "$url"
  printf 'Download destination:\n%s\n' "$destination"

  if [[ -e "$partial" ]]; then
    if [[ ! -f "$partial_url" ]] || [[ "$(<"$partial_url")" != "$url" ]]; then
      printf 'Download URL changed; restarting partial download.\n'
      rm -f -- "$partial" "$partial_url"
    else
      printf 'Resuming partial download:\n%s\n' "$partial"
    fi
  fi
  printf '%s' "$url" >"$partial_url"

  if command -v curl >/dev/null 2>&1; then
    downloader=curl
  elif command -v wget >/dev/null 2>&1; then
    downloader=wget
  else
    printf 'Error: curl or wget is required.\n' >&2
    exit 1
  fi

  if [[ "$downloader" == curl ]]; then
    if ! curl --fail --location --retry 3 --continue-at - --output "$partial" "$url"; then
      if [[ -s "$partial" ]]; then
        printf 'Resume failed; restarting download from the beginning.\n'
        rm -f -- "$partial"
        curl --fail --location --retry 3 --output "$partial" "$url"
      else
        exit 1
      fi
    fi
  else
    if ! wget --tries=3 --continue --output-document="$partial" "$url"; then
      if [[ -s "$partial" ]]; then
        printf 'Resume failed; restarting download from the beginning.\n'
        rm -f -- "$partial"
        wget --tries=3 --output-document="$partial" "$url"
      else
        exit 1
      fi
    fi
  fi

  if [[ ! -s "$partial" ]]; then
    printf 'Error: downloaded file is empty: %s\n' "$partial" >&2
    exit 1
  fi

  if [[ -n "$expected_checksum" ]] && \
      ! verify_file "$label" "$partial" "$checksum_algorithm" "$expected_checksum"; then
    rm -f -- "$partial" "$partial_url"
    return 1
  fi

  mv -- "$partial" "$destination"
  rm -f -- "$partial_url"
  printf 'Downloaded %s:\n%s\n' "$label" "$destination"
}

if wants_target debian13; then
  download_file \
    'Debian 13 cloud image' \
    "$debian_url" \
    "$debian_image" \
    sha512 \
    "$debian_sha512"
fi

if wants_target almalinux10; then
  download_file \
    'AlmaLinux 10 cloud image' \
    "$alma_url" \
    "$alma_image" \
    sha256 \
    "$alma_sha256"
fi

if wants_target fedora44; then
  download_file \
    'Fedora 44 cloud image' \
    "$fedora_url" \
    "$fedora_image" \
    sha256 \
    "$fedora_sha256"
fi

if wants_target virtio-win; then
  download_file \
    'VirtIO Windows tools ISO' \
    "$virtio_url" \
    "$virtio_image"
fi