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

ROOT_DIR="${ROOT_DIR:-/workspace}"
PROJECT_DIR="${ROOT_DIR}/reveal_vla_bimanual"
SHIM_ROOT="${ROOT_DIR}/system_shims"

export DEBIAN_FRONTEND=noninteractive

apt-get update
apt-get install -y \
  libxcb-cursor0 \
  libxcb-icccm4 \
  libxcb-image0 \
  libxcb-keysyms1 \
  libxcb-randr0 \
  libxcb-render-util0 \
  libxcb-xinerama0 \
  libxrender1 \
  libxkbcommon0 \
  libxkbcommon-x11-0 \
  mesa-utils \
  xvfb \
  x11-utils \
  x11-xserver-utils \
  xauth \
  xserver-xorg \
  xserver-xorg-core || true

if ! command -v nvidia-xconfig >/dev/null 2>&1; then
  apt-get install -y nvidia-xconfig || true
fi

mkdir -p /etc/X11/xorg.conf.d

if command -v nvidia-xconfig >/dev/null 2>&1; then
  nvidia-xconfig -a --use-display-device=None --virtual=1280x1024 || \
    nvidia-xconfig -a --virtual=1280x1024
fi

cat > /etc/X11/xorg.conf.d/99-maxclients.conf <<'EOF'
Section "ServerFlags"
    Option "MaxClients" "2048"
EndSection
EOF

cat > /etc/X11/Xwrapper.config <<'EOF'
allowed_users=anybody
needs_root_rights=yes
EOF

DRIVER_VERSION="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1)"
DRIVER_BRANCH="${DRIVER_VERSION%%.*}"
SHIM_DIR="${SHIM_ROOT}/nvidia${DRIVER_BRANCH}"

mkdir -p "${SHIM_DIR}"

XORG_DEB="$(apt-cache madison "xserver-xorg-video-nvidia-${DRIVER_BRANCH}" | awk -v v="${DRIVER_VERSION}" '$3 ~ v { print $3; exit }')"
GL_DEB="$(apt-cache madison "libnvidia-gl-${DRIVER_BRANCH}" | awk -v v="${DRIVER_VERSION}" '$3 ~ v { print $3; exit }')"

if [[ -z "${XORG_DEB}" || -z "${GL_DEB}" ]]; then
  echo "unable to resolve NVIDIA packages for driver ${DRIVER_VERSION}" >&2
  exit 1
fi

pushd "${ROOT_DIR}" >/dev/null
apt-get download \
  "xserver-xorg-video-nvidia-${DRIVER_BRANCH}=${XORG_DEB}" \
  "libnvidia-gl-${DRIVER_BRANCH}=${GL_DEB}"
dpkg-deb -x "xserver-xorg-video-nvidia-${DRIVER_BRANCH}_${XORG_DEB}_amd64.deb" "${SHIM_DIR}"
dpkg-deb -x "libnvidia-gl-${DRIVER_BRANCH}_${GL_DEB}_amd64.deb" "${SHIM_DIR}"
popd >/dev/null

cat <<'EOF'
Headless X prerequisites are installed.

Start GPU-backed X for RLBench with:
  /workspace/reveal_vla_bimanual/scripts/start_rlbench_x.sh

Then verify:
  DISPLAY=:99 glxinfo -B

The workspace-local NVIDIA Xorg shim is under:
  /workspace/system_shims
EOF