#!/bin/bash

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

set -Eeuo pipefail

bool_is_true() {
    case "${1,,}" in
        1|true|yes|on) return 0 ;;
        *) return 1 ;;
    esac
}

# Runtime and display configuration.
export DISPLAY="${DISPLAY:-:20}"
export PIPEWIRE_LATENCY="${PIPEWIRE_LATENCY:-128/48000}"
export XDG_RUNTIME_DIR="${XDG_RUNTIME_DIR:-/tmp/runtime-${USER:-desktop}}"
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}"

install -d -m 0700 "${XDG_RUNTIME_DIR}"

# Enable the optional joystick interposer only when it is actually installed.
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}"
    echo "Joystick interposer enabled: ${SELKIES_INTERPOSER_PATH}"
fi

# NVIDIA NVRTC is not always mounted by the NVIDIA container runtime.
# Keep the upstream-compatible runtime download, but place it in a
# user-writable directory instead of modifying /usr at container startup.
install_nvrtc() {
    bool_is_true "${SELKIES_DOWNLOAD_NVRTC:-true}" || return 0
    command -v nvidia-smi >/dev/null 2>&1 || return 0
    nvidia-smi >/dev/null 2>&1 || return 0

    if ldconfig -p 2>/dev/null | grep -q 'libnvrtc\.so'; then
        echo "NVRTC already available"
        return 0
    fi

    local cuda_version cuda_major nvrtc_arch multiarch
    local nvrtc_url nvrtc_archive fallback_version destination

    cuda_version="$(
        nvidia-smi --version |
        awk -F: '/CUDA Version/{gsub(/[[:space:]]/, "", $2); print $2; exit}'
    )"

    if [[ -z "${cuda_version}" ]]; then
        echo "WARNING: NVIDIA is available, but its CUDA compatibility version could not be detected" >&2
        return 0
    fi

    cuda_major="${cuda_version%%.*}"
    if [[ "${cuda_major}" =~ ^[0-9]+$ ]] && (( cuda_major >= 13 )); then
        cuda_version="12.9"
    fi

    nvrtc_arch="$(
        dpkg --print-architecture |
        sed -e 's/arm64/sbsa/' \
            -e 's/ppc64el/ppc64le/' \
            -e 's/i.*86/x86/' \
            -e 's/amd64/x86_64/' \
            -e 's/unknown/x86_64/'
    )"

    multiarch="$(
        dpkg --print-architecture |
        sed -e 's/arm64/aarch64-linux-gnu/' \
            -e 's/armhf/arm-linux-gnueabihf/' \
            -e 's/riscv64/riscv64-linux-gnu/' \
            -e 's/ppc64el/powerpc64le-linux-gnu/' \
            -e 's/s390x/s390x-linux-gnu/' \
            -e 's/i.*86/i386-linux-gnu/' \
            -e 's/amd64/x86_64-linux-gnu/' \
            -e 's/unknown/x86_64-linux-gnu/'
    )"

    nvrtc_url="https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-${nvrtc_arch}/"

    nvrtc_archive="$(
        curl -fsSL "${nvrtc_url}" |
        grep -oP "(?<=href=')cuda_nvrtc-linux-${nvrtc_arch}-${cuda_version}\.[0-9]+-archive\.tar\.xz" |
        sort -V |
        tail -n 1 ||
        true
    )"

    if [[ -z "${nvrtc_archive}" ]]; then
        fallback_version="${cuda_version}.0"
        nvrtc_archive="$(
            {
                curl -fsSL "${nvrtc_url}" |
                    grep -oP "(?<=href=')cuda_nvrtc-linux-${nvrtc_arch}-.*?\.tar\.xz" ||
                    true
                echo "cuda_nvrtc-linux-${nvrtc_arch}-${fallback_version}-archive.tar.xz"
            } |
            sort -V |
            grep -B 1 --fixed-strings "${fallback_version}" |
            head -n 1 ||
            true
        )"
    fi

    if [[ -z "${nvrtc_archive}" ]]; then
        echo "WARNING: no compatible NVRTC archive found for CUDA ${cuda_version}" >&2
        return 0
    fi

    destination="${NVRTC_DEST_PREFIX:-${XDG_CACHE_HOME:-${HOME:-/tmp}/.cache}/selkies/nvrtc}/lib/${multiarch}"
    install -d -m 0755 "${destination}"

    if compgen -G "${destination}/libnvrtc*" >/dev/null; then
        echo "Using cached NVRTC runtime from ${destination}"
        export LD_LIBRARY_PATH="${destination}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
        return 0
    fi

    echo "Installing NVRTC runtime: ${nvrtc_archive}"
    rm -rf /tmp/selkies-nvrtc
    install -d /tmp/selkies-nvrtc

    curl -fsSL "${nvrtc_url}${nvrtc_archive}" |
        tar -xJf - -C /tmp/selkies-nvrtc

    find /tmp/selkies-nvrtc -type f -name 'libnvrtc*' \
        -exec install -m 0755 {} "${destination}/" \;

    rm -rf /tmp/selkies-nvrtc
    export LD_LIBRARY_PATH="${destination}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
}

install_nvrtc

echo "Waiting for X socket ${DISPLAY}"
until [[ -S "/tmp/.X11-unix/X${DISPLAY#*:}" ]]; do
    sleep 0.5
done

echo "Waiting for PipeWire-Pulse socket ${PULSE_RUNTIME_PATH}/native"
until [[ -S "${PULSE_RUNTIME_PATH}/native" ]]; do
    sleep 0.5
done

# The current Selkies wheel contains the backend and HTML5 client.
# WebSocket transport needs only this one listening port.
export SELKIES_MODE="${SELKIES_MODE:-websockets}"
export SELKIES_ENCODER="${SELKIES_ENCODER:-h264enc}"
export SELKIES_ENABLE_RESIZE="${SELKIES_ENABLE_RESIZE:-false}"
export SELKIES_ENABLE_BASIC_AUTH="${SELKIES_ENABLE_BASIC_AUTH:-true}"
export SELKIES_BASIC_AUTH_USER="${SELKIES_BASIC_AUTH_USER:-${USER:-desktop}}"
export SELKIES_BASIC_AUTH_PASSWORD="${SELKIES_BASIC_AUTH_PASSWORD:-${PASSWD:-mypasswd}}"

echo "Starting Selkies on ${SELKIES_ADDR:-0.0.0.0}:${SELKIES_PORT:-8081}"
exec selkies \
    --addr="${SELKIES_ADDR:-0.0.0.0}" \
    --port="${SELKIES_PORT:-8081}" \
    --enable_https="false" \
    --mode="${SELKIES_MODE}" \
    --encoder="${SELKIES_ENCODER}" \
    --enable_resize="${SELKIES_ENABLE_RESIZE}" \
    --enable_basic_auth="${SELKIES_ENABLE_BASIC_AUTH}" \
    "$@"