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

PATH=/usr/sbin:/usr/bin:/sbin:/bin

HOOK_NAME=proxmenux-stack-dependencies.sh
# Longest wait for a dependency that is being backed up, and for any other lock.
BACKUP_LOCK_WAIT=${PROXMENUX_STACK_BACKUP_LOCK_WAIT:-14400}
OTHER_LOCK_WAIT=${PROXMENUX_STACK_OTHER_LOCK_WAIT:-300}
# vzdump starts again a guest that was running before its backup right after releasing it.
RESTART_GRACE=${PROXMENUX_STACK_RESTART_GRACE:-20}
POLL=${PROXMENUX_STACK_POLL:-5}
# Earlier published versions of this script, replaced on update.
PREVIOUS_RELEASES=(
  c28b69acd7cc8e66d3aec3e8b2ec86bda0e4f91193ec78549f1965f2c859c0d1
  02c50ef17f2411ef22f4d382c51b34a7a0974a446d4a155ecb63fb1e805a7c59
)

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

find_snippet_storage() {
  local storage
  if pvesm status --content snippets 2>/dev/null \
    | awk 'NR > 1 && $1 == "local" && $3 == "active" {found=1} END {exit !found}'; then
    printf 'local'
    return
  fi
  storage=$(pvesm status --content snippets 2>/dev/null \
    | awk 'NR > 1 && $3 == "active" {print $1; exit}')
  if [[ -z $storage ]]; then
    # A new Proxmox installation accepts no snippets on any storage.
    enable_local_snippets || die "No active storage supports snippets"
    storage=local
  fi
  printf '%s' "$storage"
}

enable_local_snippets() {
  local content
  content=$(pvesh get /storage/local --output-format json 2>/dev/null | jq -r '.content // empty')
  [[ -n $content ]] || return 1
  pvesm set local --content "${content},snippets" >/dev/null 2>&1 || return 1
  pvesm status --content snippets 2>/dev/null \
    | awk 'NR > 1 && $1 == "local" && $3 == "active" {found=1} END {exit !found}'
}

install_hook() {
  local main_id=${1:?missing main VMID} source_config=${2:?missing lifecycle JSON}
  local storage hook_volume hook_path target_config
  [[ $main_id =~ ^[0-9]+$ ]] || die "Invalid main VMID"
  jq -e '
    .schema == 1 and
    (.dependencies | type == "array") and
    (.dependencies | length > 0) and
    (all(.dependencies[];
      (.vmid | type == "number") and
      (.label | type == "string") and
      (.healthcheck.type | IN("exec", "http", "running")) and
      (.healthcheck.timeout_seconds | type == "number")
    ))
  ' "$source_config" >/dev/null || die "Invalid dependency contract"

  storage=$(find_snippet_storage)
  hook_volume="${storage}:snippets/${HOOK_NAME}"
  hook_path=$(pvesm path "$hook_volume")
  install -D -m 0755 "$0" "$hook_path"

  target_config="/etc/pve/priv/proxmenux-stack-${main_id}.json"
  umask 077
  cat "$source_config" >"$target_config"
  pct set "$main_id" --hookscript "$hook_volume" >/dev/null
  printf 'Proxmox hookscript installed: CT %s starts its dependencies through %s\n' \
    "$main_id" "$hook_volume"
}

# The installed copies of an earlier release take this version; a copy
# changed by hand is left as it is.
refresh_hooks() {
  local volume path current known
  while IFS= read -r volume; do
    path=$(pvesm path "$volume" 2>/dev/null) || continue
    [[ -f $path && ! -L $path ]] || continue
    current=$(sha256sum "$path" | cut -d' ' -f1)
    for known in "${PREVIOUS_RELEASES[@]}"; do
      if [[ $current == "$known" ]]; then
        install -m 0755 "$0" "$path.new" && mv -f "$path.new" "$path"
        printf 'Updated %s\n' "$path"
        break
      fi
    done
  done < <(sed -n "s/^hookscript:[[:space:]]*\([^[:space:]]*:snippets\/${HOOK_NAME//./\\.}\)[[:space:]]*$/\1/p" \
             /etc/pve/lxc/*.conf 2>/dev/null | sort -u)
}

is_running() {
  [[ $(pct status "$1" 2>/dev/null || true) == "status: running" ]]
}

container_lock() {
  pct config "$1" 2>/dev/null | sed -n '/^lock:/{s/^lock:[[:space:]]*//p;q;}' || true
}

# A dependency stopped for its own backup stays locked until vzdump has
# saved it, and vzdump starts it again when it was running before.
wait_for_release() {
  local id=$1 label=$2 lock limit waited=0
  lock=$(container_lock "$id")
  [[ -n $lock ]] || return 0
  limit=$OTHER_LOCK_WAIT
  if [[ $lock == backup ]]; then
    limit=$BACKUP_LOCK_WAIT
  fi
  printf 'Dependency %s (CT %s) is locked (%s); waiting until it is released...\n' "$label" "$id" "$lock"
  while lock=$(container_lock "$id"); [[ -n $lock ]]; do
    (( waited < limit )) || die "$label (CT $id) is still locked ($lock) after ${waited}s"
    sleep "$POLL"
    waited=$((waited + POLL))
  done
  printf 'Dependency %s (CT %s) was released.\n' "$label" "$id"
  waited=0
  while (( waited < RESTART_GRACE )) && ! is_running "$id"; do
    sleep 2
    waited=$((waited + 2))
  done
}

dependency_is_healthy() {
  local id=$1 healthcheck=$2 type url
  type=$(jq -r '.type' <<<"$healthcheck")
  case "$type" in
    running)
      [[ $(pct status "$id" 2>/dev/null || true) == "status: running" ]]
      ;;
    exec)
      local -a command=()
      mapfile -t command < <(jq -r '.argv[]' <<<"$healthcheck")
      ((${#command[@]} > 0)) || return 1
      pct exec "$id" -- "${command[@]}" >/dev/null 2>&1 9>&-
      ;;
    http)
      url=$(jq -r '.url' <<<"$healthcheck")
      curl -fsS --max-time 3 "$url" >/dev/null 2>&1
      ;;
    *) return 1 ;;
  esac
}

start_dependencies() {
  local main_id=$1 config="/etc/pve/priv/proxmenux-stack-${1}.json"
  local encoded dependency id label healthcheck timeout elapsed
  [[ -r $config ]] || die "Missing dependency contract for main CT $main_id"

  exec 9>"/run/lock/proxmenux-stack-${main_id}.lock"
  flock 9
  while IFS= read -r encoded; do
    [[ -n $encoded ]] || continue
    dependency=$(base64 -d <<<"$encoded")
    id=$(jq -r '.vmid' <<<"$dependency")
    label=$(jq -r '.label' <<<"$dependency")
    healthcheck=$(jq -c '.healthcheck' <<<"$dependency")
    timeout=$(jq -r '.healthcheck.timeout_seconds' <<<"$dependency")
    [[ $id =~ ^[0-9]+$ && $timeout =~ ^[0-9]+$ && $timeout -gt 0 ]] \
      || die "Invalid dependency in $config"
    pct config "$id" >/dev/null 2>&1 \
      || die "Dependency $label (CT $id) does not exist"

    is_running "$id" || wait_for_release "$id" "$label"
    if ! is_running "$id"; then
      printf 'Starting dependency %s (CT %s)...\n' "$label" "$id"
      # What pct start leaves running, such as the DHCP client, must not keep the lock.
      pct start "$id" 9>&- || is_running "$id" || die "Could not start $label (CT $id)"
    else
      printf 'Dependency %s (CT %s) was already running.\n' "$label" "$id"
    fi

    elapsed=0
    while (( elapsed < timeout )); do
      if dependency_is_healthy "$id" "$healthcheck"; then
        printf 'Dependency %s (CT %s): ready.\n' "$label" "$id"
        break
      fi
      [[ $(pct status "$id" 2>/dev/null || true) == "status: running" ]] \
        || die "$label (CT $id) stopped while starting"
      sleep 2
      elapsed=$((elapsed + 2))
    done
    (( elapsed < timeout )) \
      || die "$label (CT $id) did not pass its health check within ${timeout}s"
  done < <(jq -r '.dependencies[] | @base64' "$config")
}

if [[ ${1:-} == "--install" ]]; then
  shift
  install_hook "$@"
  exit 0
fi
if [[ ${1:-} == "--refresh" ]]; then
  refresh_hooks
  exit 0
fi

vmid=${1:?missing VMID}
phase=${2:?missing lifecycle phase}
case "$phase" in
  pre-start) start_dependencies "$vmid" ;;
  post-start|pre-stop|post-stop) ;;
  *) die "Unknown lifecycle phase: $phase" ;;
esac