#!/bin/bash

CACHE_PATH=/tmp/pcie40_reload.cache

#ug`pcie40_reload.description`
# This command must be executed to reload the PCIe40 driver after an FPGA has been reprogrammed.
# _ The tool performs the following steps, in order:

RELOAD_MODULE=0
RELOAD_ALL=0
USE_CACHE=0
FLUSH_CACHE=0
QUIET=0

#ug`pcie40_reload.synopsis`
# *pcie40_reload* [-m] [-a] [-c] [-f] [-q]

function usage {
  echo "pcie40_reload [-m] [-a] [-c] [-f] [-q]" >&2
  echo " -m    reload kernel module" >&2
  echo " -a    reload all devices" >&2
  echo " -c    always read cache" >&2
  echo " -f    always overwrite cache" >&2
  echo " -q    suppress dmesg output" >&2
}

while getopts "macfqh" opt; do
  case $opt in

#ug`pcie40_reload.options`mod
# *-m*::
# Remove and then re-insert module into kernel (this is required after a driver update).
    m) RELOAD_MODULE=1
      ;;
#ug`pcie40_reload.options`all
# *-a*::
# Reload all interface (by default only the interfaces where the FPGA has been reprogrammed are reloaded).
    a) RELOAD_ALL=1
      ;;
#ug`pcie40_reload.options`cache
# *-c*::
# Take the device list from the local cache, if it exists. By default the cache is used only if no PCIe40 interface appears under sysfs. With this option the cache contents are used instead. This can be useful in case a reload was issued too early (while one FPGA was still being reprogrammed) and the corresponding sysfs nodes do not exist anymore.
    c) USE_CACHE=1
      ;;
#ug`pcie40_reload.options`flush
# *-f*::
# Forces a cache update. By default the cache is updated only if additional entries have been detected compared to the cache contents. With this option the cache will be overwritten regardless.
    f) FLUSH_CACHE=1
      ;;
#ug`pcie40_reload.options`quiet
# *-q*::
# Suppress dmesg output. Without this option a partial dmesg log is printed for troubleshooting in case of errors.
    q) QUIET=1
      ;;
    h)
      usage
      exit 1
      ;;
    \?)
      echo "Invalid option: -${OPTARG}" >&2
      exit 1
      ;;
    :)
      echo "Option -${OPTARG} requires an argument" >&2
      exit 1
      ;;
  esac
done

if [ "${RELOAD_MODULE}" -eq 0 ]; then
  #INSTALLED_SRCVERSION=`modinfo -F srcversion pcie40`
  LOADED_SRCVERSION=1
  #if [ "${LOADED_SRCVERSION}" != "${INSTALLED_SRCVERSION}" ]; then
  #echo "Warning: loaded module version does not match installed version, run \`pcie40_reload -m\` to use the installed module instead" >&2
  #fi;
else
  RELOAD_ALL=1
fi;

devs=($(find /sys/devices/pci0000\:* -name pcie40_interface -exec cat {} \; -exec echo -n ' ' \; -execdir pwd \; | sort -nr | cut -f 2 -d ' ' 2>/dev/null))

if [ -z "${devs}" ]; then
  if [ -f "${CACHE_PATH}" ]; then
    echo "No PCIe40 devices appear in sysfs, using cache"
    devs=($(<${CACHE_PATH}))
  else
    echo "No devices detected! Assuming the FPGA was programmed correctly, you might have to reboot"
    exit 1
  fi
elif [ ${USE_CACHE} -eq 1 ]; then
  if [ -f "${CACHE_PATH}" ]; then
    echo "Using PCIe40 interface list from cache"
    devs=($(<${CACHE_PATH}))
  else
    echo "No PCIe40 interface cache present, if not all PCIe40 boards are detected you might have to reboot"
    exit 1
  fi
fi
if [ -f "${CACHE_PATH}" ]; then
  devs_cached=($(<${CACHE_PATH}))
  if [ ${FLUSH_CACHE} -eq 1 ] || [ ${#devs_cached[@]} -lt ${#devs[@]} ]; then
    echo "Updating cache"
    echo "${devs}" > ${CACHE_PATH}
  fi
else
  echo "Creating cache"
  echo "${devs}" > ${CACHE_PATH}
fi

ok=1

pids_ctrl=
pids_bar0=
pids_bar2=

get_pids() {
# $1: file to check
# $2: variable to append pids to
  pids=`sudo lsof -t $1`
  for pid in $pids; do
    ps -p $pid --no-headers -o "  %U%p%c"
    eval "${2}+=\ ${pid}"
  done
}

if_id=${#devs[@]}

for dev in ${devs[@]}; do
  (( if_id -= 1 ))

  if [ ${RELOAD_ALL} -eq 0 ] && [ -f "${dev}/pcie40_loaded" ] && [ $(< ${dev}/pcie40_loaded) -eq 1 ]; then
    continue;
  fi;
  echo "PCIe40 interface ${if_id}:"

  if [ $((if_id % 2)) -eq 0 ]; then
    #if [ -e /dev/pcie40_${if_id}_bar0 ]; then
    #  echo " BAR0 used by:"
    #  get_pids /dev/pcie40_${if_id}_bar0 pids_bar0
    #else
    #  echo " BAR0 missing!"
    #fi
    if [ -e /dev/pcie40_${if_id}_bar2 ]; then
      echo " BAR2 used by:"
      get_pids /dev/pcie40_${if_id}_bar2 pids_bar2
    else
      echo " BAR2 missing!"
    fi
  fi
  #if [ -e /dev/pcie40_${if_id}_ctrl ]; then
  #  echo " CTRL used by:"
  #  get_pids /dev/pcie40_${if_id}_ctrl pids_ctrl
  #else
  #  echo " CTRL missing!"
  #fi
done

#ug`pcie40_reload.steps`1
# . Terminates any process currently using the PCIe40 boards to be reloaded

if [ -n "$pids_bar0" ] || [ -n "$pids_bar2" ] || [ -n "$pids_ctrl" ]; then
  echo -n "Stopping processes... "
  kill $pids_bar0 $pids_bar2 $pids_ctrl
  echo OK
fi

#ug`pcie40_reload.steps`2
# . For every PCIe40 interface:
# .. Disconnects the interface via _sysfs_
# .. Requests a PCI rescan on the upstream device
# .. Checks if the driver detected successfully the interface

devs_removed=

if [ ${RELOAD_MODULE} -eq 1 ]; then
  sudo rmmod lhcb_pcie40 > /dev/null 2>&1
  sudo rmmod pcie40
  sudo insmod ${PCIE40_DRIVER_DIR}/pcie40.ko
else
  for dev in ${devs[@]}; do
    if [ ${RELOAD_ALL} -eq 0 ] && [ -f "${dev}/pcie40_loaded" ] && [ $(< ${dev}/pcie40_loaded) -eq 1 ]; then
      continue;
    fi;
    if [ ! -d "${dev}" ]; then
      devs_removed="${dev} ${devs_removed}"
      continue;
    fi;
    echo "Removing device ${dev}"
    if [ -f "${dev}/device" ] && [ $(< ${dev}/device) = 0xce40 ]; then
      if [ -e ${dev}/driver ]; then
        echo " Driver: OK"
      else
        echo " Driver: NO"
      fi;
      echo 1 | sudo tee ${dev}/remove > /dev/null
      devs_removed="${dev} ${devs_removed}"
    fi;
  done

  for dev in ${devs_removed}; do
    echo "Reloading device ${dev}"
    rp=`dirname ${dev}`
    echo 1 | sudo tee ${rp}/rescan > /dev/null
    if [ -e ${dev}/driver ]; then
      echo " Driver: OK"
    else
      echo " Driver: NO"
      ok=0
    fi;
  done
fi;

if_id=${#devs[@]}

for dev in ${devs[@]}; do
  (( if_id -= 1 ))
  echo "PCIe40 interface ${if_id}:"

  #ug`pcie40_reload.steps`3
  # . Checks the presence of the ECS registers

  if [ $((if_id % 2)) -eq 0 ]; then
    # echo -n " BAR0: "
    # if [ -e /dev/pcie40_${if_id}_bar0 ]; then
    #   echo "OK"
    # else
    #   echo "NO"
    #   ok=0
    # fi
    echo -n " BAR2: "
    if [ -e /dev/pcie40_${if_id}_bar2 ]; then
      echo "OK"
    else
      echo "NO"
      ok=0
    fi
  fi

  #ug`pcie40_reload.steps`4
  # . Checks the presence of the DMA controllers

  #echo -n " CTRL: "
  #if [ -e /dev/pcie40_${if_id}_ctrl ]; then
  #  echo "OK"
  #else
  #  echo "NO"
  #  ok=0
  #fi
done

if [ "$ok" -eq 0 ] && [ ${QUIET} -eq 0 ]; then
  echo "Errors found, debug output follows:"
  dmesg | grep -E "(P40:)|(P40DAQ:)" | tail -n 100
  exit 1
fi

sudo chmod a+rwx /dev/pcie40_0_*

#ug`pcie40_reload.exit`
# 0 on success, 1 if issues have been encountered. In the latter case the last diagnostic messages from the driver are also printed to standard output.