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

marker="/host/var/lib/atlas-maintenance/titan-24-docker.ready"
nvidia_volume_marker="/host/var/lib/atlas-maintenance/nvidia-driver-vol.version"
host_sh() {
  chroot /host /bin/sh -c "$*"
}

host_systemctl() {
  chroot /host /usr/bin/nsenter -t 1 -m -u -i -n -p /bin/systemctl "$@"
}

install_docker() {
  host_sh "apt-get update"
  host_sh "DEBIAN_FRONTEND=noninteractive apt-get install -y ca-certificates curl gnupg"
  host_sh "install -m 0755 -d /etc/apt/keyrings"
  host_sh "curl -fsSL https://download.docker.com/linux/debian/gpg -o /etc/apt/keyrings/docker.asc"
  host_sh "chmod a+r /etc/apt/keyrings/docker.asc"
  host_sh '. /etc/os-release; codename="${VERSION_CODENAME:-trixie}"; echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/debian ${codename} stable" > /etc/apt/sources.list.d/docker.list'
  host_sh "apt-get update"
  host_sh "DEBIAN_FRONTEND=noninteractive apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin"
}

configure_docker() {
  mkdir -p /host/etc/docker
  if [ ! -f /host/etc/docker/daemon.json ]; then
    cat >/host/etc/docker/daemon.json <<'JSON'
{
  "log-driver": "local",
  "log-opts": {
    "max-size": "50m",
    "max-file": "3"
  }
}
JSON
  fi
  if [ -x /host/usr/bin/nvidia-ctk ]; then
    host_sh "nvidia-ctk runtime configure --runtime=docker"
  fi
  host_systemctl daemon-reload
  host_systemctl enable --now docker
  host_systemctl restart docker
}

configure_nvidia_modeset() {
  if [ ! -d /host/sys/module/nvidia_drm ]; then
    return
  fi

  mkdir -p /host/etc/modprobe.d
  if ! grep -qs '^options nvidia-drm modeset=1' /host/etc/modprobe.d/nvidia-drm-modeset.conf 2>/dev/null; then
    cat >/host/etc/modprobe.d/nvidia-drm-modeset.conf <<'EOF'
options nvidia-drm modeset=1
EOF
  fi

  if [ -f /host/etc/default/grub ] && ! grep -q 'nvidia-drm.modeset=1' /host/etc/default/grub; then
    host_sh "cp /etc/default/grub /etc/default/grub.atlas-maintenance.bak"
    host_sh "sed -i 's/GRUB_CMDLINE_LINUX_DEFAULT=\"\\([^\"]*\\)\"/GRUB_CMDLINE_LINUX_DEFAULT=\"\\1 nvidia-drm.modeset=1\"/' /etc/default/grub"
    if [ -x /host/usr/sbin/update-grub ]; then
      host_sh "update-grub"
    elif [ -x /host/usr/sbin/grub-mkconfig ]; then
      host_sh "grub-mkconfig -o /boot/grub/grub.cfg"
    fi
  fi
}

ensure_nvidia_driver_volume() {
  if [ ! -r /host/sys/module/nvidia/version ]; then
    return
  fi

  local driver_version
  driver_version="$(cat /host/sys/module/nvidia/version)"

  if [ -f "${nvidia_volume_marker}" ] && grep -qx "${driver_version}" "${nvidia_volume_marker}" && host_sh "docker volume inspect nvidia-driver-vol >/dev/null 2>&1"; then
    return
  fi

  host_sh "tmp_dir=\$(mktemp -d) && curl -fsSL https://raw.githubusercontent.com/games-on-whales/gow/master/images/nvidia-driver/Dockerfile -o \"\${tmp_dir}/Dockerfile\" && docker build --build-arg NV_VERSION='${driver_version}' -t 'atlas/nvidia-driver:${driver_version}' \"\${tmp_dir}\"; rc=\$?; rm -rf \"\${tmp_dir}\"; exit \${rc}"
  host_sh "docker volume create nvidia-driver-vol >/dev/null"
  host_sh "docker run --rm -v nvidia-driver-vol:/usr/nvidia 'atlas/nvidia-driver:${driver_version}' sh -c ':'"

  mkdir -p "$(dirname "${nvidia_volume_marker}")"
  printf '%s\n' "${driver_version}" >"${nvidia_volume_marker}"
}

if [ ! -x /host/usr/bin/docker ] || ! host_sh "docker version >/dev/null 2>&1"; then
  install_docker
fi

configure_docker
configure_nvidia_modeset
ensure_nvidia_driver_volume

mkdir -p "$(dirname "${marker}")"
date -Is >"${marker}"

while true; do
  if [ ! -S /host/var/run/docker.sock ]; then
    configure_docker
    ensure_nvidia_driver_volume
  fi
  sleep 300
done