#!/bin/sh
# fillrate-task: docker-gpu-fix
# Make docker --gpus work when the NVIDIA container runtime is in cdi mode,
# then clear the failed driver-update post-reboot unit if that check passes.
#
# Mode choice: nvidia-container-toolkit 1.20.1 still accepts
# nvidia-container-runtime.mode=legacy. The v1.20.1 runtime README says only
# legacy is directly compatible with the docker --gpus flag, because that
# flag invokes the NVIDIA runtime hook. auto is the default and can select
# cdi (the default became JIT CDI in 1.18). cdi is the failure this task
# repairs, so the mode set below is legacy, not auto.
#
# Command used (1.20.1 nvidia-ctk config flags --set and --in-place;
# --config-file is the documented selector, defaulting to the same path):
#   nvidia-ctk config --set nvidia-container-runtime.mode=legacy --in-place \
#     --config-file /etc/nvidia-container-runtime/config.toml
# The docker nvidia runtime is registered only when daemon.json lacks it:
#   nvidia-ctk runtime configure --runtime=docker --config /etc/docker/daemon.json
# That command does not pass the flag that would make nvidia the default
# runtime, so an existing default runtime stays as it was.
#
# This script does not install or upgrade packages and does not change holds.
# The runner starts it with a clean environment. The FILLRATE_* path overrides
# below are for tests; a real run uses the defaults.

set -u

STATE=${FILLRATE_DOCKER_GPU_STATE:-/var/lib/fillrate-runner/docker-gpu-fix}
CONFIG=${FILLRATE_NVIDIA_CONFIG:-/etc/nvidia-container-runtime/config.toml}
CONFIG_DIR=${FILLRATE_NVIDIA_CONFIG_DIR:-/etc/nvidia-container-runtime}
CONTAINERS_SHARE=${FILLRATE_CONTAINERS_SHARE:-/usr/share/containers}
DAEMON_JSON=${FILLRATE_DAEMON_JSON:-/etc/docker/daemon.json}
DRIVER_STATE=${FILLRATE_DRIVER_STATE:-/var/lib/fillrate-runner/nvidia-driver-update}
POST_UNIT=${FILLRATE_POST_UNIT:-/etc/systemd/system/fillrate-nvidia-driver-update-post.service}
IMAGE=nvidia/cuda:12.4.1-base-ubuntu22.04
EXPIRES=${FILLRATE_EXPIRES:-0}

BEFORE_MODE=unknown
AFTER_MODE=unknown
CONFIGURED_MODE=unknown
BEFORE_GPUS=-1
GPUS_EXIT=-1
RUNTIME_EXIT=-1
ROLLED_BACK=false
DAEMON_CHANGED=false
POST_RESULT=kept
BACKED_UP=0
DOCKER_RESTARTED=0
VASTAI_WAS_ACTIVE=0
CONFIG_EXISTED=0
DAEMON_EXISTED=0
STAMP=none
BACKUP_CONFIG=
BACKUP_DAEMON=
SUM_BEFORE=absent
DEF_BEFORE=unset

refuse() {
  printf '%s\n' "refusing: $1"
  exit 1
}

window_open() {
  now=$(date +%s) || return 1
  # Same bound as the runner: the window is still open when now equals expires.
  [ "$now" -le "$EXPIRES" ]
}

read_mode() {
  if [ ! -f "$CONFIG" ] || [ -L "$CONFIG" ]; then
    printf '%s\n' unset
    return 0
  fi
  awk '
    {
      line = $0
      sub(/\r$/, "", line)
      if (line ~ /^\[/) {
        if (line ~ /^\[nvidia-container-runtime\][[:space:]]*$/) section = "runtime"
        else section = "other"
        next
      }
      if (section != "runtime") next
      if (line ~ /^[[:space:]]*#/) next
      if (line ~ /^[[:space:]]*mode[[:space:]]*=/) {
        sub(/^[[:space:]]*mode[[:space:]]*=[[:space:]]*/, "", line)
        gsub(/["'\'']/, "", line)
        sub(/[[:space:]].*$/, "", line)
        print line
        found = 1
        exit
      }
    }
    END { if (!found) print "unset" }
  ' "$CONFIG"
}

default_runtime() {
  if [ ! -f "$DAEMON_JSON" ] || [ -L "$DAEMON_JSON" ]; then
    printf '%s\n' unset
    return 0
  fi
  value=$(sed -n 's/.*"default-runtime"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' "$DAEMON_JSON" | awk 'NR==1 { print; exit }')
  if [ -z "$value" ]; then
    printf '%s\n' unset
  else
    printf '%s\n' "$value"
  fi
}

runtime_registered() {
  [ -f "$DAEMON_JSON" ] && [ ! -L "$DAEMON_JSON" ] || return 1
  grep -E -q '"nvidia"[[:space:]]*:' "$DAEMON_JSON"
}

file_sum() {
  if [ -f "$1" ] && [ ! -L "$1" ]; then
    sha256sum "$1" | awk 'NR==1 { print $1 }'
  else
    printf '%s\n' absent
  fi
}

json_word() {
  word=$(printf '%s\n' "$1" | tr -cd 'A-Za-z0-9._-')
  if [ -z "$word" ]; then
    printf '%s\n' unknown
  else
    printf '%s\n' "$word"
  fi
}

json_int() {
  case "$1" in
    -1|[0-9]|[0-9][0-9]|[0-9][0-9][0-9]) printf '%s\n' "$1" ;;
    *) printf '%s\n' -1 ;;
  esac
}

active_word() {
  if ! command -v systemctl >/dev/null 2>&1; then
    printf '%s\n' unknown
    return 0
  fi
  word=$(systemctl is-active "$1" 2>/dev/null || true)
  word=$(printf '%s\n' "$word" | awk 'NR==1 { print $1 }')
  json_word "$word"
}

filter_secrets() {
  grep -a -E -i -v \
    'nonce|ticket|secret|ssh-ed25519|private key|begin ssh signature|(^|[^A-Za-z0-9])(key|sig)([^A-Za-z0-9]|$)' \
    || true
}

print_text_file() {
  path=$1
  printf '%s\n' "=== $path ==="
  if [ -L "$path" ]; then
    printf '%s\n' "skipped symlink"
    return 0
  fi
  if [ ! -f "$path" ]; then
    printf '%s\n' "missing"
    return 0
  fi
  filter_secrets < "$path"
}

print_dir() {
  dir=$1
  printf '%s\n' "=== files under $dir ==="
  if [ ! -e "$dir" ]; then
    printf '%s\n' "directory missing"
    return 0
  fi
  if [ -L "$dir" ] || [ ! -d "$dir" ]; then
    printf '%s\n' "skipped symlink"
    return 0
  fi
  names=$(find "$dir" \( -type f -o -type l \) | sort)
  if [ -z "$names" ]; then
    printf '%s\n' "no files"
    return 0
  fi
  printf '%s\n' "$names" | while IFS= read -r path; do
    [ -n "$path" ] || continue
    base=$(basename "$path")
    if printf '%s\n' "$base" | grep -E -i -q 'ticket|key|nonce|sig|secret'; then
      printf '%s\n' "=== skip $base ==="
      continue
    fi
    print_text_file "$path"
  done
}

check_abs() {
  label=$1
  path=$2
  case "$path" in
    *..*|*' '*|*'	'*)
      refuse "bad $label path"
      ;;
    /*)
      ;;
    *)
      refuse "bad $label path"
      ;;
  esac
}

rollback() {
  printf '%s\n' "=== rollback ==="
  ROLLED_BACK=true
  if [ "$CONFIG_EXISTED" = 1 ]; then
    cp -p "$BACKUP_CONFIG" "$CONFIG" || printf '%s\n' "failed to restore config.toml"
  else
    rm -f "$CONFIG"
  fi
  if [ "$DAEMON_EXISTED" = 1 ]; then
    cp -p "$BACKUP_DAEMON" "$DAEMON_JSON" || printf '%s\n' "failed to restore daemon.json"
  else
    rm -f "$DAEMON_JSON"
  fi
  if [ "$DAEMON_CHANGED" = true ] || [ "$DOCKER_RESTARTED" = 1 ]; then
    ids=$(docker ps -q 2>/dev/null) || ids=unknown
    if printf '%s\n' "$ids" | grep -q .; then
      printf '%s\n' "not restarting docker during rollback: a container is running"
    else
      printf '%s\n' "=== systemctl restart docker ==="
      DOCKER_RESTARTED=1
      systemctl restart docker || printf '%s\n' "docker restart during rollback failed"
    fi
  fi
}

restart_docker() {
  printf '%s\n' "=== systemctl restart docker ==="
  DOCKER_RESTARTED=1
  systemctl restart docker
  code=$?
  if [ "$code" -ne 0 ]; then
    printf '%s\n' "systemctl restart docker failed ($code)"
    return 1
  fi
  if systemctl is-active --quiet docker; then
    printf '%s\n' "docker is active"
    return 0
  fi
  printf '%s\n' "docker is not active after restart"
  return 1
}

ensure_vastai() {
  if systemctl is-active --quiet vastai.service; then
    printf '%s\n' "vastai.service is active"
    return 0
  fi
  if [ "$VASTAI_WAS_ACTIVE" != 1 ]; then
    printf '%s\n' "vastai.service was not active before this task"
    return 0
  fi
  printf '%s\n' "vastai.service died; restarting"
  printf '%s\n' "=== systemctl restart vastai.service ==="
  systemctl restart vastai.service
  code=$?
  if [ "$code" -ne 0 ]; then
    printf '%s\n' "systemctl restart vastai.service failed ($code)"
    return 1
  fi
  if systemctl is-active --quiet vastai.service; then
    printf '%s\n' "vastai.service is active"
    return 0
  fi
  printf '%s\n' "vastai.service did not become active"
  return 1
}

summary() {
  AFTER_MODE=$(json_word "$(read_mode)")
  docker_state=$(active_word docker)
  vastai_state=$(active_word vastai.service)
  before_mode=$(json_word "$BEFORE_MODE")
  configured_mode=$(json_word "$CONFIGURED_MODE")
  before_gpus=$(json_int "$BEFORE_GPUS")
  gpus_exit=$(json_int "$GPUS_EXIT")
  runtime_exit=$(json_int "$RUNTIME_EXIT")
  printf '%s\n' "=== summary ==="
  cat << EOF
{
  "task": "docker-gpu-fix",
  "before_mode": "$before_mode",
  "configured_mode": "$configured_mode",
  "after_mode": "$AFTER_MODE",
  "before_gpus_exit": $before_gpus,
  "gpus_test_exit": $gpus_exit,
  "runtime_nvidia_exit": $runtime_exit,
  "docker": "$docker_state",
  "vastai": "$vastai_state",
  "post_unit": "$POST_RESULT",
  "rolled_back": $ROLLED_BACK,
  "daemon_json_changed": $DAEMON_CHANGED
}
EOF
}

fail() {
  printf '%s\n' "$1"
  if [ "$BACKED_UP" = 1 ] && [ "$ROLLED_BACK" = false ]; then
    rollback
  fi
  if [ "$DOCKER_RESTARTED" = 1 ]; then
    ensure_vastai || printf '%s\n' "vastai.service is not active after the docker change"
  fi
  summary
  exit 1
}

cleanup_post() {
  printf '%s\n' "=== cleanup fillrate-nvidia-driver-update-post ==="
  systemctl disable fillrate-nvidia-driver-update-post.service || printf '%s\n' "systemctl disable returned $?"
  systemctl reset-failed fillrate-nvidia-driver-update-post.service || printf '%s\n' "systemctl reset-failed returned $?"
  if [ -e "$POST_UNIT" ] || [ -L "$POST_UNIT" ]; then
    rm -f "$POST_UNIT" || {
      printf '%s\n' "failed to remove post unit"
      return 1
    }
  fi
  systemctl daemon-reload || printf '%s\n' "systemctl daemon-reload returned $?"
  if [ -e "$DRIVER_STATE/pending-post" ] || [ -L "$DRIVER_STATE/pending-post" ]; then
    rm -f "$DRIVER_STATE/pending-post" || {
      printf '%s\n' "failed to remove pending-post"
      return 1
    }
  fi
  if [ -e "$POST_UNIT" ] || [ -L "$POST_UNIT" ]; then
    printf '%s\n' "post unit file still exists"
    return 1
  fi
  if [ -e "$DRIVER_STATE/pending-post" ] || [ -L "$DRIVER_STATE/pending-post" ]; then
    printf '%s\n' "pending-post still exists"
    return 1
  fi
  note=$STATE/note.txt
  cat > "$note" << EOF
docker-gpu-fix: docker --gpus nvidia-smi succeeded.
Set nvidia-container-runtime.mode from $BEFORE_MODE to legacy.
Disabled and removed fillrate-nvidia-driver-update-post.service.
Removed pending-post under the driver-update state directory.
Backup stamp: $STAMP
Package holds were not changed.
EOF
  printf '%s\n' "=== note ==="
  cat "$note"
  POST_RESULT=removed
  return 0
}

check_abs STATE "$STATE"
check_abs config "$CONFIG"
check_abs "config dir" "$CONFIG_DIR"
check_abs "containers share" "$CONTAINERS_SHARE"
check_abs daemon.json "$DAEMON_JSON"
check_abs "driver state" "$DRIVER_STATE"
check_abs "post unit" "$POST_UNIT"
case "$STATE" in
  /|/etc|/usr|/var) refuse "bad state path" ;;
esac
case "$POST_UNIT" in
  /*/fillrate-nvidia-driver-update-post.service) ;;
  *) refuse "bad post unit path" ;;
esac
case "$DRIVER_STATE" in
  /) refuse "bad driver state path" ;;
esac
if [ -L "$STATE" ] || [ -L "$CONFIG" ] || [ -L "$DAEMON_JSON" ] || [ -L "$DRIVER_STATE" ]; then
  refuse "refusing to follow a symlink config or state path"
fi

if ! command -v docker >/dev/null 2>&1; then
  refuse "docker not found"
fi
ids=$(docker ps -q) || refuse "docker ps failed"
if printf '%s\n' "$ids" | grep -q .; then
  refuse "containers are running"
fi

if ! command -v nvidia-smi >/dev/null 2>&1; then
  refuse "nvidia-smi not found"
fi
if ! nvidia-smi >/dev/null 2>&1; then
  printf '%s\n' "refusing: nvidia-smi failed"
  nvidia-smi || true
  exit 1
fi

case "$EXPIRES" in
  ''|*[!0-9]*) refuse "maintenance window has passed" ;;
esac
if ! window_open; then
  refuse "maintenance window has passed"
fi

if ! command -v timeout >/dev/null 2>&1; then
  refuse "timeout not found"
fi
if ! command -v nvidia-ctk >/dev/null 2>&1; then
  refuse "nvidia-ctk not found"
fi
if ! command -v systemctl >/dev/null 2>&1; then
  refuse "systemctl not found"
fi
if ! command -v sha256sum >/dev/null 2>&1; then
  refuse "sha256sum not found"
fi
if ! command -v awk >/dev/null 2>&1; then
  refuse "awk not found"
fi

if systemctl is-active --quiet vastai.service; then
  VASTAI_WAS_ACTIVE=1
else
  VASTAI_WAS_ACTIVE=0
fi

printf '%s\n' "=== toolkit version ==="
if command -v dpkg-query >/dev/null 2>&1; then
  dpkg-query -W -f '${Package} ${Version}\n' nvidia-container-toolkit || true
else
  printf '%s\n' "dpkg-query not found"
fi

printf '%s\n' "=== nvidia-ctk --version ==="
nvidia-ctk --version || true

BEFORE_MODE=$(read_mode)
printf '%s\n' "=== before mode ==="
printf '%s\n' "before_mode=$BEFORE_MODE"

printf '%s\n' "=== config.toml ==="
print_text_file "$CONFIG"
if [ -f "$CONFIG" ] && [ ! -L "$CONFIG" ]; then
  printf '%s\n' "=== mode line ==="
  grep -n -E -C 3 '^[[:space:]]*mode[[:space:]]*=' "$CONFIG" | filter_secrets || printf '%s\n' "mode line absent"
fi

printf '%s\n' "=== daemon.json ==="
print_text_file "$DAEMON_JSON"

printf '%s\n' "=== docker info runtime ==="
docker info 2>&1 | grep -i runtime || printf '%s\n' "no runtime lines"

print_dir "$CONFIG_DIR"
print_dir "$CONTAINERS_SHARE"

printf '%s\n' "=== before: timeout 120 docker run --rm --gpus all $IMAGE nvidia-smi ==="
timeout 120 docker run --rm --gpus all "$IMAGE" nvidia-smi
BEFORE_GPUS=$?
printf '%s\n' "before_gpus_exit=$BEFORE_GPUS"

ids=$(docker ps -q) || refuse "docker ps failed"
if printf '%s\n' "$ids" | grep -q .; then
  refuse "containers are running"
fi
if ! window_open; then
  refuse "maintenance window has passed"
fi

STAMP=$(date -u +%Y%m%dT%H%M%SZ).$$
old_umask=$(umask)
umask 077
mkdir -p "$STATE" || refuse "cannot write state"
chmod 700 "$STATE" || refuse "cannot write state"
if [ -L "$STATE" ] || [ ! -d "$STATE" ]; then
  refuse "cannot write state"
fi
BACKUP_CONFIG=$STATE/config.toml.$STAMP
BACKUP_DAEMON=$STATE/daemon.json.$STAMP
if [ -f "$CONFIG" ]; then
  cp -p "$CONFIG" "$BACKUP_CONFIG" || refuse "cannot back up config.toml"
  CONFIG_EXISTED=1
else
  printf '%s\n' "config.toml absent" > "$STATE/config.toml.$STAMP.absent"
  CONFIG_EXISTED=0
fi
if [ -f "$DAEMON_JSON" ]; then
  cp -p "$DAEMON_JSON" "$BACKUP_DAEMON" || refuse "cannot back up daemon.json"
  DAEMON_EXISTED=1
else
  printf '%s\n' "daemon.json absent" > "$STATE/daemon.json.$STAMP.absent"
  DAEMON_EXISTED=0
fi
cat > "$STATE/meta.$STAMP" << EOF
before_mode=$BEFORE_MODE
config_existed=$CONFIG_EXISTED
daemon_existed=$DAEMON_EXISTED
EOF
BACKED_UP=1
umask "$old_umask"
printf '%s\n' "=== backups ==="
printf '%s\n' "$BACKUP_CONFIG"
printf '%s\n' "$BACKUP_DAEMON"

SUM_BEFORE=$(file_sum "$DAEMON_JSON")
DEF_BEFORE=$(default_runtime)

printf '%s\n' "=== nvidia-ctk config --set nvidia-container-runtime.mode=legacy --in-place ==="
nvidia-ctk config --set nvidia-container-runtime.mode=legacy --in-place --config-file "$CONFIG"
ctk_code=$?
CONFIGURED_MODE=$(read_mode)
if [ "$ctk_code" -ne 0 ] || [ "$CONFIGURED_MODE" != "legacy" ]; then
  fail "nvidia-ctk did not set mode to legacy (exit $ctk_code, mode $CONFIGURED_MODE)"
fi
printf '%s\n' "configured_mode=$CONFIGURED_MODE"

if runtime_registered; then
  printf '%s\n' "nvidia runtime already in daemon.json"
  runtime_code=0
else
  printf '%s\n' "=== nvidia-ctk runtime configure --runtime=docker ==="
  nvidia-ctk runtime configure --runtime=docker --config "$DAEMON_JSON"
  runtime_code=$?
  if [ "$runtime_code" -ne 0 ]; then
    sum_now=$(file_sum "$DAEMON_JSON")
    if [ "$sum_now" != "$SUM_BEFORE" ]; then
      DAEMON_CHANGED=true
    fi
    fail "nvidia-ctk runtime configure failed ($runtime_code)"
  fi
fi

sum_after=$(file_sum "$DAEMON_JSON")
if [ "$sum_after" != "$SUM_BEFORE" ]; then
  DAEMON_CHANGED=true
fi
def_after=$(default_runtime)
if [ "$def_after" != "$DEF_BEFORE" ]; then
  fail "default runtime changed from $DEF_BEFORE to $def_after"
fi
if ! runtime_registered; then
  fail "nvidia runtime is still missing from daemon.json"
fi

if [ "$DAEMON_CHANGED" = true ]; then
  restart_docker || fail "docker restart failed"
fi
ensure_vastai || fail "vastai.service did not become active"

printf '%s\n' "=== timeout 120 docker run --rm --gpus all $IMAGE nvidia-smi ==="
timeout 120 docker run --rm --gpus all "$IMAGE" nvidia-smi
GPUS_EXIT=$?
printf '%s\n' "gpus_test_exit=$GPUS_EXIT"

printf '%s\n' "=== docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all $IMAGE nvidia-smi ==="
docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all "$IMAGE" nvidia-smi
RUNTIME_EXIT=$?
printf '%s\n' "runtime_nvidia_exit=$RUNTIME_EXIT"

if [ "$GPUS_EXIT" -ne 0 ]; then
  fail "gpus test failed ($GPUS_EXIT); restoring backups"
fi

if ! cleanup_post; then
  printf '%s\n' "post-reboot unit cleanup failed; config changes were kept"
  summary
  exit 1
fi

summary
exit 0