gswitch

#!/bin/sh

if [ "$(whoami)" != "root" ]; then
 echo "You need to be root."
 exit
fi

OS=$(lsb_release -a 2>/dev/null | awk '/Distributor ID:/{print tolower($NF)}')

case "${OS}" in
  opensuse)
    XORG_CONF_DIR="/etc/X11/xorg.conf.d"
    if [ ! -d "${XORG_CONF_DIR}" ]; then
      mkdir -p "${XORG_CONF_DIR}"
    fi
    EGPU_TEMPLATE="${XORG_CONF_DIR}/xorg.conf.egpu"
    XORG_CONFIG="${XORG_CONF_DIR}/99-gswitch-egpu.conf"
  ;;
  ubuntu)
    XORG_CONF_DIR="/etc/X11"
    if [ ! -d "${XORG_CONF_DIR}" ]; then
      mkdir -p "${XORG_CONF_DIR}"
    fi
    EGPU_TEMPLATE="${XORG_CONF_DIR}/xorg.conf.egpu"
    XORG_CONFIG="${XORG_CONF_DIR}/xorg.conf"
  ;;
  *)
    # Untested OS distribution
    # Hoping the path works the all the same
    XORG_CONF_DIR="/etc/X11/xorg.conf.d"
    if [ ! -d "${XORG_CONF_DIR}" ]; then
      mkdir -p "${XORG_CONF_DIR}"
    fi
    EGPU_TEMPLATE="${XORG_CONF_DIR}/xorg.conf.egpu"
    XORG_CONFIG="${XORG_CONF_DIR}/99-gswitch-egpu.conf"
  ;;
esac

if [ -f "${EGPU_TEMPLATE}" ]; then
  EGPU_BUS_ID=$(awk '/BusID/{gsub(/(PCI:|\")/," ");print$2}' \
                < "${EGPU_TEMPLATE}" 2>/dev/null)
  XORG_DRIVER=$(awk '/Driver/{gsub("\"","");if($1 != "#")print$2}' \
                < "${EGPU_TEMPLATE}" 2>/dev/null)
fi
if [ "$(echo "${XORG_DRIVER}" | wc -w)" -gt "1" ]; then
  echo "More than one driver specified in the template. Exiting."
  exit 1
fi

if [ -h "${XORG_CONFIG}" ] && \
   [ -f "${EGPU_TEMPLATE}" ] && \
   [ "$(readlink "${XORG_CONFIG}")" = "${EGPU_TEMPLATE}" ]; then
  EGPU_SET="true"
else
  EGPU_SET="false"
fi

ask_reload() {
  QUESTION=$1
  echo "${QUESTION}"
  read -r ANSWER
  case "${ANSWER}" in
    Y|y|Yes|yes)
      RELOAD="true"
    ;;
    N|n|No|no)
      RELOAD="false"
    ;;
    *)
      echo "Unknown argument: ${ANSWER}."
      exit 1
    ;;
  esac
}

do_reload() {
  ( trap '' HUP TERM
    while [ "$(systemctl status display-manager | awk '/Active:/{print$2}')" \
               = "active" ]; do
      sleep 1
    done
    if [ "${REMOVE}" = "true" ]; then
      for PATH_D in ${DEV_PATH}; do
        if [ -e "${PATH_D}" ]; then
          echo 1 > "${PATH_D}"
        fi
      done
    fi
    if [ "$(lspci -k | grep -c "${XORG_DRIVER}")" -gt 0 ]; then
      modprobe "${XORG_DRIVER}"
      if [ "${XORG_DRIVER}" = "nvidia" ]; then
        modprobe nvidia_drm
      fi
    fi
    systemctl start display-manager.service ) &
  systemctl stop display-manager.service
  exit 0
}

switch_egpu() {
  if [ ! -f "${EGPU_TEMPLATE}" ]; then
    echo "No eGPU template detected, please run setup before use:"
    echo "# gswitch setup"
    exit 1
  fi
  if [ "${XORG_DRIVER}" != "nvidia" ]; then
    for HEX_ID in $(lspci | grep -i 'vga' | cut -f 1 -d ' '); do
      DEC_ID=""
      for HEX_VALUE in $(echo "${HEX_ID}" | tr ':|\.' ' '); do
        HEX_VALUE_UPPER=$(echo "${HEX_VALUE}" | tr '[:lower:]' '[:upper:]')
        DEC_VALUE=$(echo "ibase=16; ${HEX_VALUE_UPPER}" | bc)
        DEC_ID="${DEC_ID}:${DEC_VALUE}"
      done
      DEC_ID=$(echo "${DEC_ID}" | sed 's/^://')
      if [ "${DEC_ID}" = "${EGPU_BUS_ID}" ]; then
        break
      fi
    done
    DISP_PATH="/sys/bus/pci/devices/0000:${HEX_ID}/drm/card*/card*"
    DISP_STATUS=$(for DISP in ${DISP_PATH}; do find "${DISP}" -type f -name \
                  'status'; done)
    DISP_DISCON_NUM=$(printf '%s\n' "${DISP_STATUS}" | xargs \
                      grep -e '^disconnected$' | wc -l)
    DISP_TOTAL_NUM=$(printf '%s\n' "${DISP_STATUS}" | wc -l)
    if [ "${DISP_DISCON_NUM}" -eq "${DISP_TOTAL_NUM}" ] && \
       [ "${DISP_TOTAL_NUM}" -gt "0" ]; then
      echo "Warning: No eGPU attached display detected with open source drivers."
      echo "Internal mode and setting DRI_PRIME variable are recommended for this configuration."
      echo "Not setting eGPU mode."
      switch_internal
      exit 0
    fi
  fi
  if [ "${EGPU_SET}" = "true" ]; then
    ask_reload "You are already set up. Would you like to reload? (Y/n)"
    if [ "${RELOAD}" = "true" ]; then
      do_reload
    else
      exit 0
    fi
  else
    ln -s "${EGPU_TEMPLATE}" "${XORG_CONFIG}"
    ask_reload "You are now set up. Would you like to reload? (Y/n)"
    if [ "${RELOAD}" = "true" ]; then
      do_reload
    else
      exit 0
    fi
  fi
}

switch_internal() {
  if [ ! -f "${EGPU_TEMPLATE}" ]; then
    echo "No eGPU template detected, please run setup before use:"
    echo "# gswitch setup"
    exit 1
  fi
  if [ "${EGPU_SET}" = "true" ]; then
    rm -f "${XORG_CONFIG}"
    ask_reload "You are now set up. Would you like to reload? (Y/n)"
    if [ "${RELOAD}" = "true" ]; then
      do_reload
    else
      exit 0
    fi
  else
    ask_reload "You are already set up. Would you like to reload? (Y/n)"
    if [ "${RELOAD}" = "true" ]; then
      do_reload
    else
      exit 0
    fi
  fi
}

case $1 in
  -h|help)
    echo "Valid arguments are \"egpu\", \"internal\", \"boot\", \"remove\" or \"setup\"."
    echo "Example: sudo gswitch egpu"
    exit 0
  ;;
  egpu)
    switch_egpu
  ;;
  internal)
    switch_internal
  ;;
  boot)
    MODE="internal"
    for HEX_ID in $(lspci | grep -i 'vga' | cut -f 1 -d ' '); do
      DEC_ID=""
      for HEX_VALUE in $(echo "${HEX_ID}" | tr ':|\.' ' '); do
        HEX_VALUE_UPPER=$(echo "${HEX_VALUE}" | tr '[:lower:]' '[:upper:]')
        DEC_VALUE=$(echo "ibase=16; ${HEX_VALUE_UPPER}" | bc)
        DEC_ID="${DEC_ID}:${DEC_VALUE}"
      done
      DEC_ID=$(echo "${DEC_ID}" | sed 's/^://')
      if [ "${DEC_ID}" = "${EGPU_BUS_ID}" ]; then
        MODE="egpu"
        break
      fi
    done
    case ${MODE} in
      egpu)
        if [ "${EGPU_SET}" = "true" ]; then
          echo "no" | switch_egpu
        else
          echo "yes" | switch_egpu
        fi
      ;;
      internal)
        if [ "${EGPU_SET}" = "true" ]; then
          echo "yes" | switch_internal
        else
          echo "no" | switch_internal
        fi
      ;;
    esac
  ;;
  remove)
    BUS1D=$(echo "${EGPU_BUS_ID}" | cut -f 1 -d :)
    BUS2D=$(echo "${EGPU_BUS_ID}" | cut -f 2 -d :)
    BUS1H=$(printf "%02x" "${BUS1D}")
    BUS2H=$(printf "%02x" "${BUS2D}")
    HEX_DID="${BUS1H}:${BUS2H}."
    DEV_PATH="/sys/bus/pci/devices/[0-9a-f:]*${HEX_DID}[0-9]*/remove"
    echo "Try removing the GPU PCIe addresses? This will switch to internal and log you out. (Beta) (Y/n): "
    read -r REMOVE_ANSWER
    case "${REMOVE_ANSWER}" in
      Y|y|Yes|yes)
        REMOVE="true"
      ;;
      N|n|No|no)
        REMOVE="false"
      ;;
      *)
        echo "Unknown argument: ${REMOVE_ANSWER}."
        exit 1
      ;;
    esac
  echo "yes" | switch_internal
  ;;
  setup)
    cp /usr/share/gswitch/xorg.conf.egpu "${EGPU_TEMPLATE}"
    lspci | awk '/vga|VGA/{print "BusID: "$0}'
    echo "Which of these cards is your eGPU?"
    echo "Please type in the BusID, e.g: 00:02.0"
    read -r HEX_ID
    VALID_HEX_ID=$(echo "${HEX_ID}" | grep -Ec '^[a-z0-9]+:[a-z0-9]+\.[a-z0-9]$')
    if [ "${VALID_HEX_ID}" -ne 1 ]; then
      echo "You have typed in an unvalid BusID."
      echo "Example of valid entry: 00:02.0"
      exit 1
    fi
    for HEX_VALUE in $(echo "${HEX_ID}" | tr ':|\.' ' '); do
      HEX_VALUE_UPPER=$(echo "${HEX_VALUE}" | tr '[:lower:]' '[:upper:]')
      DEC_VALUE=$(echo "ibase=16; ${HEX_VALUE_UPPER}" | bc)
      DEC_ID="${DEC_ID}:${DEC_VALUE}"
    done
    DEC_ID=$(echo "${DEC_ID}" | sed 's/^:/PCI:/')
    sed -E -i "s/BusID.*$/BusID      \"${DEC_ID}\"/" "${EGPU_TEMPLATE}"
    GPU_ID=$(lspci | grep -i "${HEX_ID}")
    if echo "${GPU_ID}" | grep -qi 'nvidia'; then
      DRIVER="nvidia"
    elif echo "${GPU_ID}" | grep -qi 'amd'; then
      DRIVER="amdgpu"
    else
      echo "Couldn't automatically find the proper driver for the GPU. Exiting."
      exit 1
    fi
    sed -E -i "s/Driver.*$/Driver     \"${DRIVER}\"/" "${EGPU_TEMPLATE}"
    echo "Setup complete!"
    exit 0
  ;;
  *)
    echo "Unknown argument: ${1}."
    echo "Valid arguments are \"egpu\", \"internal\", \"boot\", \"remove\" or \"setup\"."
    echo "Example: sudo gswitch egpu"
    exit 1
  ;;
esac