#!/usr/bin/env bash

# Unloads, and removes all instances of the Graphcore IPU PCIe driver.
# The script must be run as root.

set -o nounset
set -o errexit
set -o pipefail

DRIVER_NAME="ipu_driver"

# utility functions
INFO() {
  /bin/echo -e "\e[104m\e[97m[INFO]\e[49m\e[39m $*"
}

WARNING() {
  /bin/echo >&2 -e "\e[101m\e[97m[WARNING]\e[49m\e[39m $*"
}

ERROR() {
  /bin/echo >&2 -e "\e[101m\e[97m[ERROR]\e[49m\e[39m $*"
  exit 1
}


if [ "${EUID:-999}" -ne 0 ]; then
  ERROR "Please run as root, or user with cached sudo privileges"
fi

# Unload the driver if loaded.
if modinfo ${DRIVER_NAME} >& /dev/null; then
  if ! modprobe -r ${DRIVER_NAME}; then
    WARNING "Removal of driver FAILED."
    INFO "checking for active IPU device(s)..."
    for f in /dev/ipu*[0-9];
    do
      if [ -e "$f" ]; then
        read -r process pid user <<< "$(lsof "$f" | awk '$4 ~ /mem/ { print $1 " " $2 " " $3}');"
        if [ -n "$process" ]; then
          INFO "USER-$user $process PID-$pid active on ${f:5}";
        fi
      fi
    done
    ERROR "Exiting."
    exit
  fi
fi

# Remove all old versions. 
# cope with both of these output formats from "dkms status"
# "ipu_driver, 2.0.1, 5.15.0-71-generic, x86_64: installed"
# "ipu_driver/2.0.1, 5.15.0-71-generic, x86_64: installed"
dkms status ${DRIVER_NAME} | while read line
do
  # remove first part of string: '<driver_name>, ' or '<driver_name>/'
  stripped=$(echo ${line} | sed 's/^[^ ,\/]*, //1' | sed 's/^[^ ,\/]*\///1')
  # get version string
  version=$(echo ${stripped} | grep -o '^[^,]*')
  echo "Removing ${DRIVER_NAME} version ${version}"
  INFO "$(dkms remove -m ${DRIVER_NAME} -v ${version} --all)"
done