#!/bin/bash

# Adapted from docker-selkies-glx-desktop for an XFCE-only NVIDIA image.
# This Source Code Form is subject to the Mozilla Public License, v. 2.0.

set -Eeuo pipefail

trap 'echo "Desktop entrypoint received a termination signal"' HUP INT QUIT TERM

export DISPLAY="${DISPLAY:-:20}"
export DISPLAY_SIZEW="${DISPLAY_SIZEW:-1920}"
export DISPLAY_SIZEH="${DISPLAY_SIZEH:-1080}"
export DISPLAY_REFRESH="${DISPLAY_REFRESH:-60}"
export DISPLAY_DPI="${DISPLAY_DPI:-96}"
export DISPLAY_CDEPTH="${DISPLAY_CDEPTH:-24}"
export VIDEO_PORT="${VIDEO_PORT:-DFP}"
export XDG_RUNTIME_DIR="${XDG_RUNTIME_DIR:-/tmp/runtime-${USER:-ubuntu}}"
export PIPEWIRE_RUNTIME_DIR="${PIPEWIRE_RUNTIME_DIR:-${XDG_RUNTIME_DIR}}"
export PULSE_RUNTIME_PATH="${PULSE_RUNTIME_PATH:-${XDG_RUNTIME_DIR}/pulse}"
export PULSE_SERVER="${PULSE_SERVER:-unix:${PULSE_RUNTIME_PATH}/native}"

until [[ -d "${XDG_RUNTIME_DIR}" ]]; do
    sleep 0.5
done

sudo-root chown "$(id -u):$(id -g)" "${HOME}" || true

if [[ -n "${PASSWD:-}" ]]; then
    printf '%s:%s\n' "${USER}" "${PASSWD}" | sudo-root chpasswd || true
fi

rm -rf /tmp/.X11-unix "/tmp/.X${DISPLAY#:}-lock" "${HOME}/.cache" || true

sudo-root ln -snf "/usr/share/zoneinfo/${TZ:-UTC}" /etc/localtime || true
printf '%s\n' "${TZ:-UTC}" | sudo-root tee /etc/timezone >/dev/null || true

# Joystick interposer support.
SELKIES_INTERPOSER_PATH="$(
    find /usr/lib /usr/lib64 /usr/local/lib \
        -name 'selkies_joystick_interposer.so' \
        -print -quit 2>/dev/null || true
)"

if [[ -n "${SELKIES_INTERPOSER_PATH}" ]]; then
    export LD_PRELOAD="${SELKIES_INTERPOSER_PATH}${LD_PRELOAD:+:${LD_PRELOAD}}"
    export SDL_JOYSTICK_DEVICE="${SDL_JOYSTICK_DEVICE:-/dev/input/js0}"

    sudo-root install -d -m 1777 /dev/input || true
    for joystick in /dev/input/js{0..3}; do
        sudo-root touch "${joystick}" || true
        sudo-root chmod 0777 "${joystick}" || true
    done
fi

install_nvidia_userspace() {
    command -v nvidia-smi >/dev/null 2>&1 || {
        echo "NVIDIA Container Toolkit has not exposed nvidia-smi." >&2
        return 1
    }

    if command -v nvidia-xconfig >/dev/null 2>&1; then
        return 0
    fi

    local driver_arch driver_version installer archive_dir

    driver_arch="$(
        dpkg --print-architecture |
        sed -e 's/arm64/aarch64/' \
            -e 's/armhf/32bit-ARM/' \
            -e 's/i.*86/x86/' \
            -e 's/amd64/x86_64/' \
            -e 's/unknown/x86_64/'
    )"

    driver_version="${NVIDIA_DRIVER_VERSION:-}"

    if [[ -z "${driver_version}" && -f /proc/driver/nvidia/version ]]; then
        driver_version="$(
            head -n 1 /proc/driver/nvidia/version |
            awk '{for (i=1; i<=NF; i++) if ($i ~ /^[0-9]+\.[0-9.]+$/) {print $i; exit}}'
        )"
    fi

    if [[ -z "${driver_version}" ]]; then
        driver_version="$(
            nvidia-smi --query-gpu=driver_version --format=csv,noheader |
            head -n 1
        )"
    fi

    if [[ -z "${driver_version}" ]]; then
        echo "Unable to determine the host NVIDIA driver version." >&2
        return 1
    fi

    installer="/tmp/NVIDIA-Linux-${driver_arch}-${driver_version}.run"
    archive_dir="/tmp/NVIDIA-Linux-${driver_arch}-${driver_version}"

    echo "Installing NVIDIA ${driver_version} userspace GLX components"

    curl -fsSL \
        "https://international.download.nvidia.com/XFree86/Linux-${driver_arch}/${driver_version}/NVIDIA-Linux-${driver_arch}-${driver_version}.run" \
        -o "${installer}" \
    || curl -fsSL \
        "https://international.download.nvidia.com/tesla/${driver_version}/NVIDIA-Linux-${driver_arch}-${driver_version}.run" \
        -o "${installer}"

    rm -rf "${archive_dir}"
    (
        cd /tmp
        sh "${installer}" -x
    )

    sudo-root "${archive_dir}/nvidia-installer" \
        --silent \
        --no-kernel-module \
        --no-install-compat32-libs \
        --no-nouveau-check \
        --no-nvidia-modprobe \
        --no-systemd \
        --no-rpms \
        --no-backup \
        --no-check-for-alternate-installs

    rm -rf "${installer}" "${archive_dir}"
}

install_nvidia_userspace

# Always select the first GPU visible inside the container. Device indexes are
# renumbered by NVIDIA Container Toolkit, so the host-side index is unreliable.
GPU_SELECT="$(
    nvidia-smi --query-gpu=uuid --format=csv,noheader |
    head -n 1
)"

if [[ -z "${GPU_SELECT}" ]]; then
    echo "No visible NVIDIA GPU was detected." >&2
    exit 1
fi

HEX_ID="$(
    nvidia-smi \
        --id="${GPU_SELECT}" \
        --query-gpu=pci.bus_id \
        --format=csv,noheader |
    head -n 1
)"

IFS=':.' read -r domain bus device function <<<"${HEX_ID}"

BUS_ID="PCI:$((16#${bus}))@$((16#${domain})):$((16#${device})):$((16#${function}))"
MODELINE="$(cvt -r "${DISPLAY_SIZEW}" "${DISPLAY_SIZEH}" "${DISPLAY_REFRESH}" | sed -n '2p')"
MODE_NAME="$(awk '{print $2}' <<<"${MODELINE}" | tr -d '"')"

if [[ "${VIDEO_PORT,,}" == "none" ]]; then
    CONNECTED_MONITOR=(--use-display-device=None)
else
    CONNECTED_MONITOR=(--connected-monitor="${VIDEO_PORT}")
fi

sudo-root rm -f /etc/X11/xorg.conf

sudo-root nvidia-xconfig \
    --virtual="${DISPLAY_SIZEW}x${DISPLAY_SIZEH}" \
    --depth="${DISPLAY_CDEPTH}" \
    --mode="${MODE_NAME}" \
    --allow-empty-initial-configuration \
    --no-probe-all-gpus \
    --busid="${BUS_ID}" \
    --include-implicit-metamodes \
    --mode-debug \
    --no-sli \
    --no-base-mosaic \
    --only-one-x-screen \
    "${CONNECTED_MONITOR[@]}"

sudo-root sed -i \
    '/Driver[[:space:]]*"nvidia"/a\    Option "ModeValidation" "NoMaxPClkCheck,NoEdidMaxPClkCheck,NoMaxSizeCheck,NoHorizSyncCheck,NoVertRefreshCheck,NoVirtualSizeCheck,NoExtendedGpuCapabilitiesCheck,NoTotalSizeCheck,NoDualLinkDVICheck,NoDisplayPortBandwidthCheck,AllowNon3DVisionModes,AllowNonHDMI3DModes,AllowNonEdidModes,NoEdidHDMI2Check,AllowDpInterlaced"' \
    /etc/X11/xorg.conf

sudo-root sed -i \
    '/Driver[[:space:]]*"nvidia"/a\    Option "AllowExternalGpus" "True"' \
    /etc/X11/xorg.conf

sudo-root sed -i \
    '/Section[[:space:]]*"Monitor"/a\    '"${MODELINE}" \
    /etc/X11/xorg.conf

sudo-root sed -i '/"DPMS"/d' /etc/X11/xorg.conf
sudo-root sed -i \
    '/Section[[:space:]]*"Monitor"/a\    Option "DPMS" "False"' \
    /etc/X11/xorg.conf

cat <<'EOF' | sudo-root tee -a /etc/X11/xorg.conf >/dev/null
Section "ServerFlags"
    Option "DontVTSwitch" "True"
    Option "DontZap" "True"
    Option "AllowMouseOpenFail" "True"
    Option "AutoAddGPU" "False"
EndSection
EOF

sudo-root ln -snf /dev/ptmx /dev/tty7

/usr/lib/xorg/Xorg \
    "${DISPLAY}" \
    vt7 \
    -noreset \
    -novtswitch \
    -sharevts \
    -nolisten tcp \
    -ac \
    -dpi "${DISPLAY_DPI}" \
    +extension COMPOSITE \
    +extension DAMAGE \
    +extension GLX \
    +extension RANDR \
    +extension RENDER \
    +extension MIT-SHM \
    +extension XFIXES \
    +extension XTEST \
    &

XORG_PID=$!

echo "Waiting for X socket ${DISPLAY}"
until [[ -S "/tmp/.X11-unix/X${DISPLAY#*:}" ]]; do
    if ! kill -0 "${XORG_PID}" 2>/dev/null; then
        echo "Xorg exited before creating its socket." >&2
        exit 1
    fi
    sleep 0.5
done

xset -dpms || true
xset s off || true
xset s noblank || true

export XDG_SESSION_ID="${DISPLAY#*:}"

xfce4-session &
XFCE_PID=$!

wait "${XFCE_PID}"