#!/bin/bash

set -eu

. "${SNAP}/usr/share/nvidia-container-toolkit/lib"

# Just exit if NVIDIA support is disabled #
[ "${NVIDIA_SUPPORT_DISABLED}" != "true" ] || exit 0

NVIDIA_ENABLE="false"

# Determine if we can enable NVIDIA support or not #
snapctl is-connected graphics-core22 && NVIDIA_ENABLE="true"
[ "${NVIDIA_SUPPORT_CLASSIC}" == "true" ] && NVIDIA_ENABLE="true"

# Ensure nvidia support setup correctly, and only if hardware is preset and correct #
if [ "${NVIDIA_ENABLE}" == "true" ] ; then

    [ "${NVIDIA_SUPPORT_CLASSIC}" == "true" ] && echo "Running on Classic system" || echo "Running on Ubuntu Core system"

    # Connection hooks are run early - copy the config file from $SNAP into $SNAP_DATA if it doesn't exist
    if [ ! -f "$SNAP_DATA/config/daemon.json" ]; then
        mkdir -p "$SNAP_DATA/config"
        cp "$SNAP/config/daemon.json" "$SNAP_DATA/config/daemon.json"
    fi

    # Ensure hardware present #
    nvidia_hw_ensure

    # As service order is not guaranteed outside of snap - wait a bit for nvidia assemble to complete #
    device_wait /dev/nvidiactl || exit 0
    echo "NVIDIA ready"

    # Ensure the data dirs exist #
    ensure_nvidia_data_dirs

    # Setup nvidia support, but do not exit on failure #
    cdi_generate && nvidia_runtime_config && docker_runtime_configure && setup_info || setup_fail

fi