#!/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"
  ubuntu_image="${base_qcow2_dir}/ubuntu-26.04-server-cloudimg-arm64-20260927.img"
  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
    https://mirrors.bfsu.edu.cn/fedora/releases/44/Cloud/aarch64/images/Fedora-Cloud-Base-Generic-44-1.7.aarch64.qcow2
    https://mirrors.nju.edu.cn/fedora/releases/44/Cloud/aarch64/images/Fedora-Cloud-Base-Generic-44-1.7.aarch64.qcow2"
  ubuntu_url_default="https://mirrors.tuna.tsinghua.edu.cn/ubuntu-cloud-images/resolute/20260927/resolute-server-cloudimg-arm64.img
    https://mirrors.nju.edu.cn/ubuntu-cloud-images/releases/26.04/release-20260927/ubuntu-26.04-server-cloudimg-arm64.img
    https://cloud-images.ubuntu.com/releases/26.04/release-20260927/ubuntu-26.04-server-cloudimg-arm64.img"
  ubuntu_sha256_default="63a93bd5a8d76e33b15ceb5daa3657bd79be804748051ab178e643b0f5da22e7"
  # 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"
  ubuntu_image="${base_qcow2_dir}/ubuntu-26.04-server-cloudimg-amd64-20260927.img"
  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
    https://mirrors.bfsu.edu.cn/fedora/releases/44/Cloud/x86_64/images/Fedora-Cloud-Base-Generic-44-1.7.x86_64.qcow2
    https://mirrors.nju.edu.cn/fedora/releases/44/Cloud/x86_64/images/Fedora-Cloud-Base-Generic-44-1.7.x86_64.qcow2"
  ubuntu_url_default="https://mirrors.tuna.tsinghua.edu.cn/ubuntu-cloud-images/resolute/20260927/resolute-server-cloudimg-amd64.img
    https://mirrors.nju.edu.cn/ubuntu-cloud-images/releases/26.04/release-20260927/ubuntu-26.04-server-cloudimg-amd64.img
    https://cloud-images.ubuntu.com/releases/26.04/release-20260927/ubuntu-26.04-server-cloudimg-amd64.img"
  debian_sha512_default="f6978100d8031c266d55d7815ceea7fcdeacf28e1e5834fdb9c94ac96880a054a6e6f8681c2d3b0584e0057eaf3ef7353856b85212d04134744faa9b3bb1f24f"
  alma_sha256_default="47f2218668dd4776be140dd92fa3bea700be1766e2c7d88bdfd6a4b50f477b4d"
  fedora_sha256_default="28680fe5b371a5a82ebf43a31926e086a168e59949d03969c5093e7071f90b7f"
  ubuntu_sha256_default="8800651811af9a85465ad1d552add729947bb16488dddb4a9b5305a3d97332b2"
fi

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

# Win32-OpenSSH names its zip after the architecture: Win64 for x64 and ARM64
# for arm64.
if [[ "$image_arch" == arm64 ]]; then
  openssh_zip=OpenSSH-ARM64.zip
else
  openssh_zip=OpenSSH-Win64.zip
fi
openssh_image="${tools_iso_dir}/${openssh_zip}"

# Windows 11 media is named after the architecture: x64 on amd64 hosts and
# Arm64 on arm64 hosts.
if [[ "$image_arch" == arm64 ]]; then
  windows_arch=arm64
else
  windows_arch=x64
fi

# Find a Windows 11 file for the current architecture in a directory by
# matching "Win", "11", the architecture 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*"${windows_arch}"*."${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 %s ISO (*Win*11*%s*.iso) found in %s\n' \
      "$windows_arch" "$windows_arch" "$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 %s ISO (*Win*11*%s*.iso) found in %s\n' \
      "$windows_arch" "$windows_arch" "$install_iso_dir" >&2
    return 1
  fi
}

# Each *_url may list several mirrors of the same file, separated by
# whitespace. aria2c downloads from all of them at once; curl and wget try
# them in order.
debian_url="${DEBIAN_IMAGE_URL:-$debian_url_default}"
alma_url="${ALMA_IMAGE_URL:-$alma_url_default}"
fedora_url="${FEDORA_IMAGE_URL:-$fedora_url_default}"
ubuntu_url="${UBUNTU_IMAGE_URL:-$ubuntu_url_default}"
openssh_url="${OPENSSH_WIN_ZIP_URL:-https://github.com/PowerShell/Win32-OpenSSH/releases/download/10.0.0.0p2-Preview/${openssh_zip}}"
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}"
ubuntu_sha256="${UBUNTU_IMAGE_SHA256:-$ubuntu_sha256_default}"
force_download=false

usage() {
  cat <<EOF
Usage: $(basename "$0") [--force] [all|debian13|almalinux10|fedora44|ubuntu2604|virtio-win|openssh-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/Ubuntu image architecture (arm64 or amd64)
is auto-detected from the host; override it with VM_IMAGE_ARCH.

When aria2c is installed and an image has several mirrors, the image is
downloaded from all of them at once, with one connection per mirror.
Otherwise curl or wget downloads from the mirrors in order.

--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 (*Win*11*arm64*.iso on arm64); 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" ;;
    ubuntu2604) printf '%s\n' "$ubuntu_image" ;;
    virtio-win) printf '%s\n' "$virtio_image" ;;
    openssh-win) printf '%s\n' "$openssh_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|ubuntu2604|virtio-win|openssh-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 to $partial with curl or wget from one URL, resuming an existing
# partial file and restarting from the beginning if resuming fails.
download_single_source() {
  local downloader="$1"
  local url="$2"
  local partial="$3"

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

# Download to $partial with aria2c, splitting the file across all URLs with a
# single connection to each server.
download_multi_source() {
  local partial="$1"
  shift

  aria2c \
    --continue=true \
    --split="$#" \
    --max-connection-per-server=1 \
    --file-allocation=none \
    --auto-file-renaming=false \
    --max-tries=5 \
    --retry-wait=3 \
    --summary-interval=0 \
    --console-log-level=warn \
    --dir="$(dirname -- "$partial")" \
    --out="$(basename -- "$partial")" \
    "$@"
}

download_file() {
  local label="$1"
  local urls
  read -r -d '' -a urls <<<"$2" || true
  local destination="$3"
  local checksum_algorithm="${4:-}"
  local expected_checksum="${5:-}"
  local partial="${destination}.part"
  # Records how the partial file was started, so that a partial download is
  # only resumed with the same downloader and URLs.
  local partial_url="${partial}.url"
  local downloader
  local download_state
  local url

  if ((${#urls[@]} == 0)); then
    printf 'Error: no download URL for %s.\n' "$label" >&2
    exit 1
  fi

  if [[ "$force_download" == true ]]; then
    printf 'Forcing a fresh download of %s.\n' "$label"
    rm -f -- "$destination" "$partial" "${partial}.aria2" "$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

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

  printf 'Downloading %s with %s from:\n' "$label" "$downloader"
  printf '%s\n' "${urls[@]}"
  printf 'Download destination:\n%s\n' "$destination"

  # aria2c fills a partial file out of order, so it must not be resumed by
  # curl or wget, and vice versa.
  download_state="$(printf '%s\n' "$downloader" "${urls[@]}")"
  if [[ -e "$partial" ]]; then
    if [[ ! -f "$partial_url" ]] || [[ "$(<"$partial_url")" != "$download_state" ]]; then
      printf 'Download source changed; restarting partial download.\n'
      rm -f -- "$partial" "${partial}.aria2" "$partial_url"
    else
      printf 'Resuming partial download:\n%s\n' "$partial"
    fi
  fi
  printf '%s\n' "$download_state" >"$partial_url"

  if [[ "$downloader" == aria2c ]]; then
    if ! download_multi_source "$partial" "${urls[@]}"; then
      printf 'Error: failed to download %s.\n' "$label" >&2
      exit 1
    fi
  else
    for url in "${urls[@]}"; do
      # Some mirrors reject clients with an empty "200 OK" response.
      if download_single_source "$downloader" "$url" "$partial" && [[ -s "$partial" ]]; then
        break
      fi
      if [[ "$url" == "${urls[-1]}" ]]; then
        printf 'Error: failed to download %s.\n' "$label" >&2
        exit 1
      fi
      printf 'Download from %s failed; trying the next mirror.\n' "$url" >&2
    done
  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}.aria2" "$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 ubuntu2604; then
  download_file \
    'Ubuntu 26.04 cloud image' \
    "$ubuntu_url" \
    "$ubuntu_image" \
    sha256 \
    "$ubuntu_sha256"
fi

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

if wants_target openssh-win; then
  download_file \
    'Win32-OpenSSH zip' \
    "$openssh_url" \
    "$openssh_image"
fi