#!/bin/bash
set -Eeuo pipefail

export DISPLAY="${DISPLAY:-:20}"
export DISPLAY_SIZEW="${DISPLAY_SIZEW:-1920}"
export DISPLAY_SIZEH="${DISPLAY_SIZEH:-1080}"
export DISPLAY_REFRESH="${DISPLAY_REFRESH:-120}"
export DISPLAY_DPI="${DISPLAY_DPI:-96}"
export DISPLAY_CDEPTH="${DISPLAY_CDEPTH:-24}"
export VIDEO_PORT="${VIDEO_PORT:-HDMI-0}"

export XDG_RUNTIME_DIR="${XDG_RUNTIME_DIR:-/tmp/runtime-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}"

XORG_PID=""
XFCE_PID=""

cleanup() {
    echo "[desktop] Stopping XFCE/Xorg..."

    if [[ -n "${XFCE_PID}" ]] && kill -0 "${XFCE_PID}" 2>/dev/null; then
        kill -TERM "${XFCE_PID}" 2>/dev/null || true
    fi

    if [[ -n "${XORG_PID}" ]] && kill -0 "${XORG_PID}" 2>/dev/null; then
        kill -TERM "${XORG_PID}" 2>/dev/null || true
    fi
}

trap cleanup EXIT HUP INT QUIT TERM

echo "[desktop] Preparing X11 runtime..."

sudo install -d -m 1777 /tmp/.X11-unix /tmp/.ICE-unix
sudo rm -f "/tmp/.X${DISPLAY#:}-lock"

# Do not wipe the whole user cache: Steam may use it.
rm -rf "${HOME}/.cache/sessions" 2>/dev/null || true

if ! command -v nvidia-smi >/dev/null 2>&1; then
    echo "[desktop] ERROR: nvidia-smi is not available." >&2
    exit 1
fi

if ! command -v nvidia-xconfig >/dev/null 2>&1; then
    echo "[desktop] ERROR: nvidia-xconfig is not available." >&2
    exit 1
fi

GPU_SELECT="$(
    nvidia-smi         --query-gpu=uuid         --format=csv,noheader     | head -n 1     | tr -d '[:space:]'
)"

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

HEX_ID="$(
    nvidia-smi         --id="${GPU_SELECT}"         --query-gpu=pci.bus_id         --format=csv,noheader     | head -n 1     | tr -d '[:space:]'
)"

if [[ -z "${HEX_ID}" ]]; then
    echo "[desktop] ERROR: unable to determine NVIDIA PCI bus ID." >&2
    exit 1
fi

IFS=':.' read -r DOMAIN BUS DEVICE FUNCTION <<<"${HEX_ID}"

BUS_ID="PCI:$((16#${BUS}))@$((16#${DOMAIN})):$((16#${DEVICE})):$((16#${FUNCTION}))"

echo "[desktop] GPU UUID : ${GPU_SELECT}"
echo "[desktop] GPU PCI  : ${HEX_ID} -> ${BUS_ID}"

MODELINE="$(
    cvt -r         "${DISPLAY_SIZEW}"         "${DISPLAY_SIZEH}"         "${DISPLAY_REFRESH}"     | sed -n '2p'
)"

if [[ -z "${MODELINE}" ]]; then
    echo "[desktop] ERROR: cvt failed to generate a modeline." >&2
    exit 1
fi

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

echo "[desktop] Generating /etc/X11/xorg.conf for ${DISPLAY_SIZEW}x${DISPLAY_SIZEH}@${DISPLAY_REFRESH}..."

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

sudo 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 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 sed -i     '/Driver[[:space:]]*"nvidia"/a\    Option         "AllowExternalGpus" "True"'     /etc/X11/xorg.conf

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

sudo sed -i '/"DPMS"/d' /etc/X11/xorg.conf

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

cat <<'EOF' | sudo tee -a /etc/X11/xorg.conf >/dev/null

Section "ServerFlags"
    Option         "DontVTSwitch" "True"
    Option         "DontZap" "True"
    Option         "AllowMouseOpenFail" "True"
    Option         "AutoAddGPU" "False"
EndSection
EOF

# Xorg is started with -sharevts exactly like the previously validated stack.
sudo ln -snf /dev/ptmx /dev/tty7

echo "[desktop] Starting Xorg on ${DISPLAY}..."

# Important:
# Do not use udevadm trigger here. Sunshine input hotplug is handled by the
# controlled Xorg restart in sunshine-entrypoint.sh.
/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 "[desktop] Waiting for X socket..."

for _ in $(seq 1 120); do
    if [[ -S "/tmp/.X11-unix/X${DISPLAY#:}" ]]; then
        break
    fi

    if ! kill -0 "${XORG_PID}" 2>/dev/null; then
        echo "[desktop] ERROR: Xorg exited before creating its socket." >&2
        wait "${XORG_PID}" || true
        exit 1
    fi

    sleep 0.5
done

if [[ ! -S "/tmp/.X11-unix/X${DISPLAY#:}" ]]; then
    echo "[desktop] ERROR: timed out waiting for Xorg." >&2
    exit 1
fi

echo "[desktop] Xorg is ready."

xset -display "${DISPLAY}" -dpms 2>/dev/null || true
xset -display "${DISPLAY}" s off 2>/dev/null || true
xset -display "${DISPLAY}" s noblank 2>/dev/null || true

export XDG_SESSION_ID="${DISPLAY#:}"
export DESKTOP_SESSION=xfce
export XDG_SESSION_DESKTOP=xfce
export XDG_CURRENT_DESKTOP=XFCE
export XDG_SESSION_TYPE=x11

echo "[desktop] Starting XFCE..."

xfce4-session &
XFCE_PID=$!

# Keep the Supervisor program alive for the lifetime of the desktop session.
wait "${XFCE_PID}"