# Configure the Nvidia driver
_configure-nvidia ACTION="prompt":
    #!/usr/bin/bash
    source /usr/lib/ujust/ujust.sh
    OPTION={{ ACTION }}
    if [ "$OPTION" == "prompt" ]; then
      echo "${bold}Configuring Nvidia drivers${normal}"
      echo 'What would you like to do?'
      OPTION=$(ugum choose "Set needed kernel arguments" "Remove kernel arguments")
    elif [ "$OPTION" == "help" ]; then
      echo "Usage: ujust configure-nvidia <option>"
      echo "  <option>: Specify the quick option - 'kargs', 'test-cuda' or 'firefox-vaapi'"
      echo "  Use 'kargs' to Set needed kernel arguments"
      echo "  Use 'test-cuda' to Test CUDA support"
      echo "  Use 'firefox-vaapi' to Enable Nvidia VAAPI in Firefox Flatpak."
      exit 0
    fi
    if [ "$OPTION" == "Set needed kernel arguments" ] || [ "${OPTION,,}" == "kargs" ]; then
        rpm-ostree kargs \
          --append-if-missing=rd.driver.blacklist=nouveau \
          --append-if-missing=modprobe.blacklist=nouveau \
          --append-if-missing=nvidia-drm.modeset=1 \
          --delete-if-present=nomodeset
    elif [ "$OPTION" == "Remove kernel arguments" ] || [ "${OPTION,,}" == "remove-kargs" ]; then
        rpm-ostree kargs \
          --delete-if-present=rd.driver.blacklist=nouveau \
          --delete-if-present=modprobe.blacklist=nouveau \
          --delete-if-present=nvidia-drm.modeset=1 \
          --append-if-missing=nomodeset
    fi

# Switch between Nvidia image and NVK
toggle-nvk:
    #!/usr/bin/bash
    CURRENT_IMAGE=$(rpm-ostree status -b --json | jq -r '.deployments[0]."container-image-reference"' | sed -E 's/^.+\/(.+:.+)$/\1/')
    CURRENT_URI=$(rpm-ostree status -b --json | jq -r '.deployments[0]."container-image-reference"' | sed -E 's/^(.+\/).+:.+$/\1/')
    NEW_IMAGE=$CURRENT_IMAGE
    CONFIG_ACTION=""
    if [[ "$CURRENT_IMAGE" =~ "nvidia" ]]; then
      NEW_IMAGE=${CURRENT_IMAGE/-nvidia/}
      CONFIG_ACTION="kargs"
    else
      NEW_IMAGE=${CURRENT_IMAGE/:/-nvidia:}
      CONFIG_ACTION="remove-kargs"
    fi
    echo "Rebasing to ${NEW_IMAGE}"
    rpm-ostree rebase ${CURRENT_URI}${NEW_IMAGE}
    just _configure-nvidia ${CONFIG_ACTION}