#!/bin/sh

set -eu

images_dir=/images

# use a common lock with puavo-bootserver-sync-images
# (only one instance of these two programs should be running)
image_updates_lock_path="${images_dir}/.image_updates_lock"
test -e "$image_updates_lock_path" || touch "$image_updates_lock_path"
exec 3< "$image_updates_lock_path"
if ! flock -nx 3; then
  echo did not get a lock on $image_updates_lock_path, exiting
  exit 1
fi

exec 4>&1       # another way to write to stdout for log()
exec 5>&2       # another way to write to stderr for log()

#
# argument parsing
#

log() {
  local logmessage logpriority log_to_stderr log_to_stdout

  log_to_stdout=false
  log_to_stderr=false

  logpriority=$1
  logmessage=$2

  case "$logpriority" in
    warn|err|crit|alert|emerg)
      log_to_stderr=true
      ;;
    *)
      $quiet && return 0
      log_to_stdout=true
      ;;
  esac

  $log_to_stdout && echo "$logmessage" >&4
  $log_to_stderr && echo "$logmessage" >&5

  echo "$logmessage" \
    | logger -p "$logpriority" -t puavo-install-and-update-ltspimages
}

usage() {
  cat <<EOF > /dev/stderr
Usage:
  $(basename $0) [--quiet] [--install-from-file path] [next_ltsp_image_name]"
  $(basename $0) [--quiet] [--install-from-nbd path]  next_ltsp_image_name"
  $(basename $0) [--quiet] [--rate-limit rate]        next_ltsp_image_name"

  -r / --rate-limit must match regexp [0-9]+[km] for kilobytes/megabytes per second
EOF
  exit 1
}

if ! args=$(getopt -n "$0" -o +f:n:r:q \
                   -l 'hosttype:,install-from-file:,install-from-nbd:,rate-limit:,quiet,images-dir:,no-preinst-hook,use-image-torrents,use-rdiff-torrents' \
                   -- "$@"); then
  usage
fi

incoming_images_dir="${images_dir}/incoming"
metas_dir="${images_dir}/meta"
rdiffs_dir="${images_dir}/rdiffs"
puavoinstall_libdir=/usr/lib/puavo-ltsp-install

hosttype=""
image_from_file=""
image_from_nbd=""
rate_limit=""
quiet=false
run_preinst_hook=true
use_torrents_for_image_updates=false
use_torrents_for_rdiff_updates=false

eval "set -- $args"
while [ $# -ne 0 ]; do
  case "$1" in
    --hosttype)
      hosttype="$2"; shift; shift
      ;;
    --images-dir)
      images_dir="$2"; shift; shift
      ;;
    --no-preinst-hook)
      run_preinst_hook=false; shift
      ;;
    -f|--install-from-file)
      image_from_file="$2"; shift; shift
      ;;
    -n|--install-from-nbd)
      image_from_nbd="$2"; shift; shift
      ;;
    -r|--rate-limit)
      # Acceptable rate limit parameter must be something that
      # wget (--limit-rate) understands.
      rate_limit="$2"; shift; shift
      if ! echo "$rate_limit" | grep -Eqx '[0-9]+[km]'; then
        log err "the rate limit was not understood, got '$rate_limit'"
        usage
      fi
      ;;
    -q|--quiet)
      quiet=true; shift
      ;;
    --use-image-torrents)
      use_torrents_for_image_updates=true; shift
      ;;
    --use-rdiff-torrents)
      use_torrents_for_rdiff_updates=true; shift
      ;;
    --) shift; break;;
  esac
done

next_image=${1:-}
if [ -z "$next_image" ]; then
  if [ -n "$image_from_file" ]; then
    next_image=$(basename "$image_from_file")
  else
    usage
  fi
fi

if [ -z "${UKI_MODE:-}" ]; then
  PUAVO_INSTALL_LIBDIR="/usr/lib/puavo-ltsp-install"
  export UKI_MODE
  UKI_MODE=$("$PUAVO_INSTALL_LIBDIR"/is-uki-install "$images_dir")
fi

#
# functions related to exiting
#

update_stats() {
  local phase progress stats_filepath

  phase=$1
  progress=$2

  stats_filepath="${images_dir}/image_update.stats"

  {
    echo "image ${next_image}" # $next_image is "global", just use it here
    echo "phase ${phase}"
    echo "progress ${progress}"
  } | tr -d '\r' > "${stats_filepath}.tmp"
    # tr -d here is nice because progress may contain \r ...
    # perhaps originates from unbuffer(?) and should be stripped earlier

  if ! cmp "${stats_filepath}" "${stats_filepath}.tmp" >/dev/null 2>&1; then
    mv "${stats_filepath}.tmp" "${stats_filepath}"
    test -x "${puavoinstall_libdir}/update-ltspimage-update-stats" \
         && "${puavoinstall_libdir}/update-ltspimage-update-stats" \
               "${stats_filepath}" 2>/dev/null || true
  else
    rm -f "${stats_filepath}.tmp"
  fi
}

interrupted=false

finish() {
  if [ "$?" -ne 0 ]; then
    if $interrupted; then
      update_stats interrupted 0
    else
      update_stats error 0
    fi
  fi
}

interrupted() {
  interrupted=true
  exit 1
}

trap finish 0
trap interrupted INT TERM


#
# subroutines
#

check_checksums_exist_for() {
  local file_to_check_checksum_for

  for file_to_check_checksum_for in "$@"; do
    if ! lookup_from_metas sha512 "$file_to_check_checksum_for" >/dev/null; then
      log err "could not find sha512 for '$file_to_check_checksum_for'"
      return 1
    fi
  done

  return 0
}

check_image_and_put_it_to_use() {
  local next_image next_image_incoming_path next_image_incoming_tmppath \
        next_image_path next_image_tmppath
  next_image=$1
  next_image_incoming_tmppath=$2

  check_checksum "$next_image" "$next_image_incoming_tmppath" image_checksum \
    || return 1

  next_image_incoming_path="${incoming_images_dir}/${next_image}"
  next_image_path="${images_dir}/${next_image}"
  next_image_tmppath="${next_image_path}.tmp"

  if ! sync; then
    log err 'failed to run sync before putting image into use'
    return 1
  fi
  if ! chown root:root "$next_image_incoming_tmppath"; then
    log err "failed to chown ${next_image_incoming_tmppath}"
    return 1
  fi
  if ! mv "$next_image_incoming_tmppath" "$next_image_incoming_path"; then
    log err "failed to move ${next_image_incoming_tmppath}"
    return 1
  fi
  if ! ln -f "$next_image_incoming_path" "$next_image_tmppath"; then
    log err "failed to hardlink ${next_image_incoming_path}"
    return 1
  fi
  if ! mv "$next_image_tmppath" "$next_image_path"; then
    log err "failed to movie ${next_image_tmppath} to ${next_image_path}"
    return 1
  fi
  if ! sync; then
    log err 'failed to run sync after putting image into use'
    return 1
  fi

  log notice "new ltsp image ${next_image} has been put to use"
}

check_checksum() {
  local actual_file_path calculated_checksum expected_checksum filename phase

  filename=$1
  actual_file_path=$2
  phase=$3

  if ! [ -r "$actual_file_path" ]; then
    log err "could not find a file ${filename} to check"
    return 1
  fi

  expected_checksum=$(lookup_from_metas sha512 "$filename")
  if [ -z "$expected_checksum" ]; then
    log err "could not determine expected checksum for ${filename}"
    return 1
  fi

  calculated_checksum=$(sha512sum_with_progress "$phase" "$actual_file_path")
  if [ -z "$calculated_checksum" ]; then
    log err "could not calculate checksum for ${filename}"
    return 1
  fi

  if [ "$expected_checksum" != "$calculated_checksum" ]; then
    log err "file ${filename} failed checksum check, removing ${actual_file_path}"
    rm -f "$actual_file_path"
    return 1
  fi

  log info "checksum for file ${filename} is okay"
  return 0
}

sha512sum_with_progress() {
  local phase sha512_filepath

  phase=$1
  sha512_filepath=$2

  # exchange stdout and stderr, redirect stderr to stdout after that
  {
    pv -n "$sha512_filepath" 3>&1 1>&2 2>&3 3>&- \
      | update_stats_with_progress "$phase"
  } 2>&1 | sha512sum | awk '{ print $1 }'
}

cleanup_previous_runs() {
  local old_sync_files preserve_imagefile preserve_rdifffile

  preserve_imagefile=$1
  preserve_rdifffile=$2

  rm -f ${images_dir}/*.img.tmp ${images_dir}/*.backup

  # cleanup old images (and their possible temporary *.tmp files)
  # and rdiff-files too.  The ltsp-backup.img is going to go too...
  old_sync_files=$("${puavoinstall_libdir}/ls-old-images-and-rdiffs" \
                     "$images_dir" "$rdiffs_dir" "$incoming_images_dir")
  if [ -n "$old_sync_files" ]; then
    printf %s "$old_sync_files" \
      | fgrep -vx "${images_dir}/${preserve_imagefile}"              \
      | fgrep -vx "${incoming_images_dir}/${preserve_imagefile}"     \
      | fgrep -vx "${incoming_images_dir}/${preserve_imagefile}.tmp" \
      | fgrep -vx "${rdiffs_dir}/${preserve_rdifffile}"              \
      | fgrep -vx "${rdiffs_dir}/${preserve_rdifffile}.tmp"          \
      | xargs -r -d'\n' -L1 rm -f
  fi

  if [ "${UKI_MODE:-}" = "1" ]; then
    with_systemd_inhibit puavo-cleanup-uki-installations \
      --images-dir "$images_dir" || true
  fi

  # Ask puavo-image-torrent-updated and puavo-rdiff-torrent-updated
  # to stop seeding (if they are doing anything).  This is relevant here
  # only for disk space calculations.
  if $use_torrents_for_image_updates; then
    echo stop | nc -N -U /run/puavo/puavo-image-torrent-updated.sock \
      >/dev/null 2>&1 || true
  fi
  if $use_torrents_for_rdiff_updates; then
    echo stop | nc -N -U /run/puavo/puavo-rdiff-torrent-updated.sock \
      >/dev/null 2>&1 || true
  fi
}

# "stat -c %s" is not good enough for files with holes and that matters here
lookup_allocated_space() {
  local values
  values=$(stat -c '%b %B' -- "$1" 2>/dev/null) || return 1
  set -- $values
  printf "%s\n" "$(($1 * $2))" || return 1
}

ensure_enough_available_diskspace() {
  local diskspace msg next_image rdiff_filename required_diskspace \
        target_image_size target_rdiff_size target_rdiff_size_on_fs \
        target_sizes_summed target_sizes_on_fs target_tmprdiff_size_on_fs \
        target_tmpimage_size_on_fs

  next_image=$1
  rdiff_filename=$2

  if ! target_image_size=$(lookup_from_metas size "$next_image"); then
    log err "could not look up target image size for $next_image"
    return 1
  fi
  target_tmpimage_size_on_fs=$(
    lookup_allocated_space "${incoming_images_dir}/${next_image}.tmp" \
      || echo 0)

  if [ -z "$rdiff_filename" ] \
    || ! target_rdiff_size=$(lookup_from_metas size "$rdiff_filename"); then
      # We are not using rdiffs for updates.
      target_rdiff_size=0
      target_rdiff_size_on_fs=0
      target_tmprdiff_size_on_fs=0
  elif [ -e "${rdiffs_dir}/${rdiff_filename}" ]; then
    # This implies rdiff should be ok.
    target_rdiff_size_on_fs="$target_rdiff_size"
    target_tmprdiff_size_on_fs=0
  else
    # No rdiff on filesystem, check out temporary rdiff size.
    target_rdiff_size_on_fs=0
    target_tmprdiff_size_on_fs=$(
      lookup_allocated_space "${rdiffs_dir}/${rdiff_filename}.tmp" || echo 0)
  fi

  target_sizes_summed=$(echo "${target_image_size} + ${target_rdiff_size}" \
                          | bc -l)
  target_sizes_on_fs=$(
    echo "${target_tmpimage_size_on_fs} + ${target_rdiff_size_on_fs} \
      + ${target_tmprdiff_size_on_fs}" | bc -l)

  required_diskspace=$(
    echo "${target_sizes_summed} - ${target_sizes_on_fs}" | bc -l)

  if ! diskspace=$(/usr/lib/puavo-ltsp-client/puavo-check-available-space \
                     --sync "$images_dir"); then
    log err 'could not determine available diskspace'
    return 1
  fi

  msg="checking if there is enough available diskspace:"
  msg="$msg diskspace:${diskspace} >= required_diskspace:$required_diskspace"
  msg="$msg = targets:${target_sizes_summed} - onfs:${target_sizes_on_fs}"
  msg="$msg = (image:${target_image_size} + rdiff:${target_rdiff_size})"
  msg="$msg - (tmpimage_on_fs:${target_tmpimage_size_on_fs} + rdiff_on_fs:${target_rdiff_size_on_fs} + tmprdiff_on_fs:${target_tmprdiff_size_on_fs})"

  log info "$msg"

  # This check presumes rdiffs directory is in the same partition as images
  # directory, sharing the same diskspace.
  if [ "$diskspace" -lt "$required_diskspace" ]; then
    msg="only $diskspace bytes available on $images_dir,"
    msg="$msg need $required_diskspace bytes"
    log err "$msg"
    return 1
  fi
}

ensure_that_default_image_is_the_current_one() {
  local booted_image default_image

  booted_image=$(cat /etc/puavo-image/name || true)
  default_image=$(lookup_ltspimage_name_by_alias ltsp.img || true)

  if [ -z "$booted_image" ]; then
    log err "could not determine the current ltspimage"
    return 1
  fi

  if [ "$booted_image" != "$default_image" ]; then
    if [ ! -e "${images_dir}/${booted_image}" ]; then
      log err "booted from image '${booted_image}', but it does not exist!"
      return 1
    fi

    log notice \
        "not booted the default image, setting '${booted_image}' as default"
    set_image_as_default_image "$booted_image"
  fi
}

fetch_with_wget() {
  local as_someone output_path phase rate_limit url wget wget_error_code \
        wgetopts

  as_someone=$1
  rate_limit=$2
  output_path=$3
  url=$4
  phase=$5

  # rdiffs may be fetched with torrents, running as "puavo-update".
  # To avoid permission problems (and maybe even for security),
  # fetch rdiffs always as "puavo-update" user.
  case "$as_someone" in
    as_puavo_update) wget='sudo -n -u puavo-update wget' ;;
    as_root)         wget='wget'                         ;;
    *)
      log err 'Internal error, fetch_with_wget called with bad arguments' >&2
      return 1
      ;;
  esac

  wgetopts=""
  if [ -n "$rate_limit" ]; then
    wgetopts="--limit-rate=$rate_limit"
  fi

  wget_error_code=$(
    {
      {
        $wget --ca-certificate=/etc/puavo-conf/rootca.pem        \
              --certificate=/etc/puavo/certs/hostorgcabundle.pem \
              --continue                                         \
              --output-document="$output_path"                   \
              --private-key=/etc/puavo/certs/host.key            \
              --progress=dot:mega                                \
              $wgetopts                                          \
              "$url" >/dev/null                                  \
          || echo $?
      } 3>&1 1>&2 2>&3 3>&- \
        | unbuffer -p awk '
            BEGIN { progress = -1 }

            {
              if (match($8, /^([[:digit:]]+)%$/, a)) {
                new_progress = a[1]
                if (progress != new_progress) {
                  progress = new_progress
                  print progress
                }
              }
            }
          ' 2>/dev/null \
        | update_stats_with_progress "$phase"
    } 2>&1 || true
  )

  if [ -n "$wget_error_code" ]; then
    log err "fetching ${url}, wget error code ${wget_error_code}"
    return 1
  fi

  log info "fetched ${url} with success"
}

get_meta_files() {
  local gnupg_dir image_server meta_download_ok meta_file series_url \
        series_url_list series_url_with_gpg_suffix

  rate_limit=$1

  meta_download_ok=false

  if ! series_url_list=$(puavo-conf puavo.image.series.urls | jq -r '.[]'); then
    log err 'could not determine series urls, is puavo.image.series.urls okay?'
    return 1
  fi
  if [ -z "$series_url_list" ]; then
    log err 'no series urls have been set'
    return 1
  fi

  mkdir -p "$metas_dir"
  for series_url in $series_url_list; do
    series_url_with_gpg_suffix="${series_url}.gpg"
    meta_file="${metas_dir}/${series_url##*/}"
    if ! fetch_with_wget as_root "$rate_limit"                 \
                                 "${meta_file}.gpg.tmp"        \
                                 "$series_url_with_gpg_suffix" \
                                 checksums_fetch; then
      log warn "failed to fetch series ${series_url}"
      rm -f "${meta_file}.gpg.tmp"
      continue
    fi

    image_server=${series_url_with_gpg_suffix#*://} # remove scheme
    image_server=${image_server%%/*}                # remove path
    gnupg_dir="/etc/puavo-conf/gpg/${image_server}/os-images"
    if ! gpg --decrypt --homedir "$gnupg_dir" --no-permission-warning \
      --trust-model always "${meta_file}.gpg.tmp" 2>/dev/null \
      > "${meta_file}.tmp"; then
        log warn "GPG verification failed for ${series_url_with_gpg_suffix}"
        rm -f "${meta_file}.gpg.tmp" "${meta_file}.tmp"
        continue
    fi
    log info "GPG verification OK for ${series_url_with_gpg_suffix}"
    rm -f "${meta_file}.gpg.tmp"

    if ! mv "${meta_file}.tmp" "$meta_file"; then
      log warn "failed to move ${series_url} file into its place"
      continue
    fi
    meta_download_ok=true
  done

  # We return success if at least one meta file download was successful.
  # This means that a single bad URL will not block updates, but a warning
  # is still shown.
  if ! $meta_download_ok; then
    return 1
  fi

  return 0
}

get_next_image() {
  local errmsg httpurl image_server image_server_list next_image rate_limit \
        previous_image rdiff_filename rdiff_should_exist

  next_image=$1
  rate_limit=$2

  rdiff_should_exist=false

  image_server_list=$(lookup_image_servers)

  previous_image=$(lookup_ltspimage_name_by_alias ltsp.img || true)

  if [ -z "$previous_image" ]; then
    log err 'could not determine the current ltsp image name'
    return 1
  fi

  if ! rdiff_filename=$(get_rdiff_filename $previous_image $next_image); then
    log err "failed to determine the rdiff filename between '${previous_image}' and '${next_image}'"
    return 1
  fi

  install -d -D -o root -g puavo-update -m 775 "$incoming_images_dir"
  install -d -D -o root -g puavo-update -m 775 "$rdiffs_dir"

  cleanup_previous_runs "${next_image}.tmp" "$rdiff_filename"

  get_meta_files "$rate_limit" || return 1

  for image_server in $image_server_list; do
    if try_rdiff_update_from_imageserver "$image_server"   \
                                         "$rate_limit"     \
                                         "$rdiff_filename" \
                                         "$previous_image" \
                                         "$next_image"; then
      # start seeding torrent (optional, it is okay if it fails)
      httpurl="https://${image_server}/${next_image}"
      timeout -k 2 5 printf "seed\n%s\n%s\n" "$httpurl" "$rdiffs_dir" \
        | nc -N -U "/run/puavo/puavo-image-torrent-updated.sock" \
          || true
      return 0
    elif [ $? -ne 2 ]; then
      # try_rdiff_update_from_imageserver returns 2 if some required file
      # does not exist in the catalogue, meaning that rdiff updates will fail.
      # Other error exit codes imply that an rdiff should exist.
      rdiff_should_exist=true
    fi
  done

  # Temporary errors (network outage etc.) should not make us fall back to
  # full image updates, because those take so much more bandwidth.
  if $rdiff_should_exist; then
    errmsg="rdiff updates failed but rdiff should exist,"
    errmsg="${errmsg} not falling back to trying a full image update"
    log err "$errmsg"
    return 1
  fi

  # We have failed with updates through rdiff, we try full image update
  # from all servers in $image_server_list.

  for image_server in $image_server_list; do
    httpurl="https://${image_server}/${next_image}"
    try_full_image_update_from_imageserver "$httpurl"    \
                                           "$rate_limit" \
                                           "$next_image" \
      && return 0
  done

  return 1
}

download_url_with_torrent() {
  local phase rate_limit torrent_type url_to_download_with_torrent

  rate_limit=$1
  url_to_download_with_torrent=$2
  torrent_type=$3
  phase=$4

  # torrent_type must be either "image" or "rdiff"
  torrent_status=$(
    printf "%s\n%s\n" "$rate_limit" "$url_to_download_with_torrent" \
      | nc -N -U "/run/puavo/puavo-${torrent_type}-torrent-updated.sock" \
      | while read progress; do
          if [ "$progress" = 'OK' ]; then
            echo OK
          else
            update_stats "$phase" "$progress"
          fi
        done)
  if [ "$torrent_status" != 'OK' ]; then
    log err "could not download ${url_to_download_with_torrent} with torrent"
    return 1
  fi

  log info "got ${url_to_download_with_torrent} with success"

  return 0
}

get_image() {
  local httpurl next_image_tmppath rate_limit

  rate_limit=$1
  next_image_tmppath=$2
  httpurl=$3

  if $use_torrents_for_rdiff_updates; then
    download_url_with_torrent "$rate_limit" "$httpurl" image image_download \
      || return 1
  else
    log info "fetching ${httpurl} with wget"
    fetch_with_wget as_puavo_update "$rate_limit"         \
                                    "$next_image_tmppath" \
                                    "$httpurl"            \
                                    image_download        \
      || return 1
  fi
}

get_rdiff() {
  local rate_limit rdiff_filename rdiff_path rdiff_tmp rdiff_url

  rate_limit=$1
  rdiff_filename=$2
  rdiff_url=$3

  rdiff_path="${rdiffs_dir}/${rdiff_filename}"

  if [ -e "$rdiff_path" ]; then
    log info "the rdiff file ${rdiff_filename} is already in place"
    return 0
  fi

  log info "we are missing the full rdiff ${rdiff_filename}, going to get it"

  rdiff_tmp="${rdiff_path}.tmp"

  if $use_torrents_for_rdiff_updates; then
    download_url_with_torrent "$rate_limit" "$rdiff_url" rdiff rdiff_fetch \
      || return 1
  else
    fetch_with_wget as_puavo_update "$rate_limit" "$rdiff_tmp" "$rdiff_url" \
      rdiff_fetch || return 1
  fi

  check_checksum "$rdiff_filename" "$rdiff_tmp" rdiff_checksum \
    || return 1

  sync || return 1
  mv "$rdiff_tmp" "$rdiff_path" || return 1

  log notice "new rdiff file ${rdiff_filename} has been put into its place"
}

get_rdiff_filename() {
  local previous_image_name next_image_name

  previous_image_name=$1
  next_image_name=$2

  echo "$previous_image_name $next_image_name" \
    | awk '
        NR == 1 {
          orig   = $1
          target = $2
          regex  = "^(.*?)-([0-9]{4}-[0-9]{2}-[0-9]{2}-[0-9]{6})-(.*?).img$"

          if (match(orig, regex, orig_match) \
            && match(target, regex, target_match)) {
              printf("%s-%s--%s-%s.rdiff\n",
                     orig_match[1],
                     orig_match[2],
                     target_match[2],
                     orig_match[3])
              exit(0)
          }
          else { exit(1) }
        }
      '
}

install_from_file_or_nbd() {
  local imagename imagepath image_src_path preserve_imagefile total_size \
        type_arg type_opt

  imagename=$1
  type_opt=$2
  type_arg=$3

  imagepath="${images_dir}/${imagename}"

  test -e "$imagepath" && return 0

  case "$type_opt" in
    -file)
      image_src_path=$type_arg
      preserve_imagefile=$(basename "$image_src_path")
      total_size=$(du -k "$image_src_path" | awk '{ print $1 "k" }')
      ;;
    -nbd)
      image_src_path=$type_arg
      preserve_imagefile=''
      total_size=$(df "$image_src_path" \
                     | awk -v image_src_path="$image_src_path" '
                         $1 == image_src_path { print $2 "k" }
                       ')
      ;;
    *)
      log err "internal error in install_from_file_or_nbd()"
      return 1
      ;;
  esac

  cleanup_previous_runs "$preserve_imagefile" ''

  dd "if=${image_src_path}" 2>/dev/null | pv -s "$total_size" \
    > "${imagepath}.tmp"
  sync
  mv "${imagepath}.tmp" "$imagepath"
  sync
}

install_image() {
  local current_ltspimage_name imagename type_arg type_opt

  imagename=$1
  type_opt=$2
  type_arg=$3

  current_ltspimage_name=$(lookup_ltspimage_name_by_alias ltsp.img || true)
  if [ "$current_ltspimage_name" = "$imagename" ]; then
    echo "The image ${imagename} is already in place."
    update_stats uptodate 100
    return 0
  fi

  if [ -n "$current_ltspimage_name" ]; then
    query_update_confirmation "$current_ltspimage_name" "$imagename"
  fi

  install_from_file_or_nbd   "$imagename" "$type_opt" "$type_arg"
  install_uki_files          "$imagename"
  run_image_preinst_hook     "$imagename"
  set_image_as_default_image "$imagename"

  echo "The image ${imagename} is now installed and set as default."
}

lookup_from_metas() {
  local meta_field meta_filename meta_value

  meta_field=$1
  meta_filename=$2

  meta_value=$(
    case "$meta_filename" in
      *.img)
        jq -r --arg field "$meta_field" --arg filename "$meta_filename" \
          '.[][][] | select(has("filename") and has($field)
                              and .filename == $filename)
                   | .[$field]' "$metas_dir"/*.json | head -1
        ;;
      *.rdiff)
        jq -r --arg field "$meta_field" --arg filename "$meta_filename" \
          '.[][][] | select(has("diffs")) | .diffs[]
                   | select(has("filename") and has($field)
                              and .filename == $filename)
                   | .[$field]' "$metas_dir"/*.json | head -1
        ;;
    esac)

  if [ -z "$meta_value" ]; then
    log err "could not lookup field $meta_field for $meta_filename from metas"
    return 1
  fi

  printf "%s\n" "$meta_value"
}

# this is mostly reimplemented in Ruby in puavo-pkg-updater
lookup_image_servers() {
  local all_image_servers image_server_by_dns image_servers_by_puavoconf \
        toplevel_image_server

  all_image_servers=''

  if image_server_by_dns=$(/usr/lib/puavo-ltsp-client/lookup-image-server-by-dns); then
    all_image_servers="$image_server_by_dns"
  else
    log info 'could not find the image server from DNS'
  fi

  # Add some extra image servers in case those have been set up in puavo-conf
  # (can be empty, one address, or many separated by whitespace).
  image_servers_by_puavoconf=$(puavo-conf puavo.image.servers)
  if [ -n "$image_servers_by_puavoconf" ]; then
    all_image_servers="${all_image_servers} ${image_servers_by_puavoconf}"
  fi

  toplevel_image_server="cdn.$(cat /etc/puavo/topdomain)"
  all_image_servers="${all_image_servers} ${toplevel_image_server}"

  log info "using image servers: $(echo "$all_image_servers" | xargs)"

  printf %s "$all_image_servers"
}

lookup_ltspimage_name_by_alias() {
  local current_ltspimage_inode current_ltspimage_path file imagename

  imagename=$1

  current_ltspimage_path="${images_dir}/${imagename}"
  current_ltspimage_inode="$(stat -c %i "$current_ltspimage_path" 2>/dev/null \
                               || true)"

  for file in ${images_dir}/*.img; do
    # check that *.img expands to something
    test -e "$file" || continue

    # ltsp.img is not what we are looking for
    test "$file" = "$current_ltspimage_path" && continue

    # we want its other name...
    if [ "$(stat -c %i "$file" || true)" = "$current_ltspimage_inode" ]; then
      echo "$(basename "$file")"
      return 0
    fi
  done

  return 1
}

patch_with_rdiff() {
  local next_image next_image_filesize next_image_tmppath previous_image \
        rdiff_filename

  previous_image=$1
  next_image=$2
  rdiff_filename=$3

  if ! next_image_filesize=$(lookup_from_metas size "$next_image"); then
    log err "could not lookup image size from metas for '${next_image}'"
    return 1
  fi

  next_image_tmppath="${incoming_images_dir}/${next_image}.tmp"
  install -o puavo-update -g puavo-update -m 644 /dev/null "$next_image_tmppath"

  # rdiff might fail due to a corrupt rdiff-file or for some other reason.
  # We let rdiff pass through in case of failure, and we check the correctness
  # of the output right after (removing the output if the checksum is bad).
  # (Note that we do not use pipefail, so with pv the rdiff status code does
  # not matter anyway).
  {
    rdiff patch "${images_dir}/${previous_image}" \
                "${rdiffs_dir}/${rdiff_filename}" \
                -                                 \
      | { pv -n -s "$next_image_filesize" 3>&1 1>&2 2>&3 3>&- \
            | update_stats_with_progress image_patch; } \
      > "$next_image_tmppath" 2>&1
  } || true

  if ! check_image_and_put_it_to_use "$next_image"         \
                                     "$next_image_tmppath"; then
    log err "error putting patched image ${next_image} into use"
    return 1
  fi

  return 0
}

query_update_confirmation() {
  local current_ltspimage_name new_ltspimage_name

  current_ltspimage_name=$1
  new_ltspimage_name=$2

  cat <<EOF
The current PuavoLTSP image is ${current_ltspimage_name},
but version ${new_ltspimage_name} is available.

Press ENTER to proceed to update it.
EOF
  read answer
}

with_systemd_inhibit() {
  systemd-inhibit --what=shutdown:sleep \
                  --who=puavo-install-and-update-ltspimages \
                  --why='Updating system' \
                  --mode=block \
    "$@"
}

install_uki_files() {
  local imagename imagepath install_hosttype

  imagename="$1"
  imagepath="${images_dir}/${imagename}"

  [ "${UKI_MODE:-}" = "1" ] || return 0

  # use the specified host type or find out the active one
  install_hosttype="$hosttype"

  if [ -z "$install_hosttype" ]; then
    install_hosttype=$(cat /etc/puavo/hosttype 2>/dev/null || true)

    if [ -z "$install_hosttype" ]; then
      log err "failed to determine host type for UKI installation"
      return 1
    fi
  fi

  with_systemd_inhibit puavo-cleanup-uki-installations \
    --images-dir "$images_dir" || true

  log info "installing UKI files from ${imagepath}"
  with_systemd_inhibit puavo-install-and-update-uki-image \
    "$imagepath" "$install_hosttype" "$images_dir"
}

run_image_preinst_hook() {
  local imagename imagepath

  if ! $run_preinst_hook; then
    return 0
  fi

  imagename=$1
  imagepath="${images_dir}/${imagename}"

  # this should do at least grub configuration update

  mkdir -p "${images_dir}/mnt"
  umount -f "${images_dir}/mnt" 2>/dev/null || true
  mount -o ro "$imagepath" "${images_dir}/mnt"

  if [ -n "$hosttype" ]; then
    "${images_dir}/mnt/${puavoinstall_libdir}/puavo-image-preinst" \
      --hosttype "$hosttype"
  else
    "${images_dir}/mnt/${puavoinstall_libdir}/puavo-image-preinst"
  fi

  umount "${images_dir}/mnt" || true
}

set_image_as_default_image() {
  local backup_ltspimage_path default_ltspimage_path imagename imagepath

  imagename=$1
  imagepath="${images_dir}/${imagename}"

  backup_ltspimage_path="${images_dir}/ltsp-backup.img"
  default_ltspimage_path="${images_dir}/ltsp.img"

  ln -f "$default_ltspimage_path" "$backup_ltspimage_path" 2>/dev/null || true
  ln -f "$imagepath" "${default_ltspimage_path}.tmp"
  sync
  mv "${default_ltspimage_path}.tmp" "$default_ltspimage_path"
  sync

  update_image_labels_for_grub
  sync

  log notice "new ltsp image $next_image has been set as default"

  update_stats finished 100
}

try_full_image_update_from_imageserver() {
  local httpurl next_image rate_limit

  httpurl=$1
  rate_limit=$2
  next_image=$3

  # it is not sensible to continue if the required checksum does not exist
  check_checksums_exist_for "$next_image" || return 1

  ensure_enough_available_diskspace "$next_image" '' || return 1

  next_image_tmppath="${incoming_images_dir}/${next_image}.tmp"
  get_image "$rate_limit" "$next_image_tmppath" "$httpurl" \
    || return 1

  if ! check_image_and_put_it_to_use "$next_image"         \
                                     "$next_image_tmppath"; then
    log err "error putting downloaded image ${next_image} into use"
    return 1
  fi

  return 0
}

try_rdiff_update_from_imageserver() {
  local image_server next_image previous_image rate_limit rdiff_filename \
        rdiff_url

  image_server=$1
  rate_limit=$2
  rdiff_filename=$3
  previous_image=$4
  next_image=$5

  # it is not sensible to continue if the required checksums do not exist
  # error 2 is special code for caller: this can not possibly work.
  check_checksums_exist_for "$rdiff_filename" "$next_image" \
    || return 2

  ensure_enough_available_diskspace "$next_image" "$rdiff_filename" || return 1

  rdiff_url="https://${image_server}/rdiffs/${rdiff_filename}"

  get_rdiff "$rate_limit"     \
            "$rdiff_filename" \
            "$rdiff_url"      \
    || return 1

  patch_with_rdiff "$previous_image" \
                   "$next_image"     \
                   "$rdiff_filename" \
    || return 1
}

update_image() {
  local next_image rate_limit

  next_image=$1
  rate_limit=$2

  if [ ! -e "${images_dir}/${next_image}" ]; then
    ensure_that_default_image_is_the_current_one

    log info "we are missing $next_image, going to get it"
    get_next_image "$next_image" "$rate_limit"
  fi

  # must call /usr/bin/test because /bin/sh is broken regarding this test
  if /usr/bin/test "${images_dir}/${next_image}" \
               -ef "${images_dir}/ltsp.img"; then
    log info "the requested ltsp image is already in place"
    update_stats uptodate 100
  else
    install_uki_files          "$next_image"
    run_image_preinst_hook     "$next_image"
    set_image_as_default_image "$next_image"
  fi
}

update_image_labels_for_grub() {
  local backup_image default_image

  default_image=$(lookup_ltspimage_name_by_alias ltsp.img || true)
  backup_image=$(lookup_ltspimage_name_by_alias ltsp-backup.img || true)

  rm -f ${images_dir}/*.backup ${images_dir}/*.default

  test -n "$default_image" \
    && touch "${images_dir}/${default_image%.img}.default"
  test -n "$backup_image"  \
    && touch "${images_dir}/${backup_image%.img}.backup"

  # This setup_grub_environment script updates the
  # "puavo_grub_puavo_os_default_image" and "puavo_grub_puavo_os_backup_image"
  # Grub environment variables to match the above default/backup images.
  /etc/puavo-conf/scripts/setup_grub_environment --images-dir "$images_dir"
}

update_stats_with_progress() {
  local phase

  phase=$1

  # possible phases are:
  #   starting
  #   checksums_fetch
  #   rdiff_fetch
  #   rdiff_checksum
  #   image_patch
  #   image_download
  #   image_checksum
  #   finished
  #   uptodate

  while read progress; do
    update_stats "$phase" "$progress"
  done
}

#
# main
#

# be_very nice
ionice -c 3  -p $$
renice -n 20 -p $$ > /dev/null

# do nothing if /images/DISABLE_IMAGE_UPDATES exists
# (for development and emergencies)
if [ -e "${images_dir}/DISABLE_IMAGE_UPDATES" ]; then
  msg="not updating image because ${images_dir}/DISABLE_IMAGE_UPDATES exists"
  log notice "$msg"
  exit 0
fi

update_stats starting 0

if [ -n "$image_from_file" ]; then
  install_image "$next_image" -file "$image_from_file"
elif [ -n "$image_from_nbd" ]; then
  install_image "$next_image" -nbd  "$image_from_nbd"
else
  update_image "$next_image" "$rate_limit"
fi

exit 0
