#!/bin/bash

show_help() {
cat << EOF
Usage: salpack-rebuild-md [--force|-f] [--watch|-w] [-a] [--status|-s]
   or: salpack-rebuild-md --clone source_sdx destination_sdx [-W]
         - clone partition table and add partitions
         - use -W for write-mostly
   or: salpack-rebuild-md [--mirror|--add] source_sdx destination_sdx [-W]
         - add partitions from source disk to destination disk
   or: salpack-rebuild-md [--write-mostly|-W] [sdXY|vdXY] ...
Use -a for both --force & --watch.
EOF
exit 0
}

SYSDIR=/sys/devices/virtual/block
MDS="`echo $SYSDIR/md* | sed s:$SYSDIR/::g`"

getparam() {
  cat $SYSDIR/$MD/md/$1
}

getdevs() {
  echo $SYSDIR/$MD/md/dev-* | sed s:$SYSDIR/$MD/md/dev-::g
}

hasstate() {
  grep -q "$2" $SYSDIR/$MD/md/dev-$1/state > /dev/null
}

echo_hr() {
  echo "---------------------------------------------------------------------"
}

echo_run() {
  echo "$@"
  $@
}

clone_parts() {
  src="$1"
  tgt="$2"
  for dev in /dev/$tgt[0-9]*; do
    if [ -b $dev ]; then
      echo "ERROR: Patition $dev found! Unable to continue."
      echo "Use wipefs on target device before continue to force."
      exit 1
    fi
  done
  set -e  # exit on any error
  echo "Cloning /dev/$src to /dev/$tgt ..."
  echo_hr
  fdisk -l /dev/$src
  echo_hr
  fdisk -l /dev/$tgt
  echo_hr
  read -p "Press enter to continue or ctrl+c to quit ..." read_enter
  echo "sfdisk -d /dev/$src | sfdisk /dev/$tgt"
  sfdisk -d /dev/$src | sfdisk /dev/$tgt
  # clone gpt sdX to sdY and randimize guid
  # sgdisk /dev/sdX -R /dev/sdY; sgdisk -G /dev/sdY
}

clone_md() {
  src="$1"
  tgt="$2"
  args="$3"
  for MD in $MDS; do
    if [ -d $SYSDIR/$MD/md ]; then
      devs=`getdevs`
      for dev in $devs; do
        tgtdev="${dev/#$src/$tgt}"
        if [ "$tgtdev" = "$dev" ]; then
          echo "Unable to map $dev to $tgt. Ignored."
        else
          echo_run "mdadm /dev/$MD -a $args /dev/$tgtdev"
        fi
      done
    fi
  done
  echo_hr
  cat /proc/mdstat
}

update_write_mostly() {
  if [ -L $SYSDIR/md*/slaves/$1 ]; then
    md=`echo $SYSDIR/md*/slaves/$1 | cut -d/ -f6`
    mdadm /dev/$md --fail /dev/$1 --remove /dev/$1 --add -W /dev/$1
    cat /proc/mdstat
  else
    echo "ERROR: $1 not found for any md device."
    ls -1d $SYSDIR/md*/slaves/*
  fi
}

md_status() {
  sed -e 's| (.*/.*)||' -e 's|[0-9]\{3\}K/sec|M/s|' "$@"
}
export -f md_status

for param in "$@"; do
  if [ "$param" = "--help" -o "$param" = "-h" ]; then
    show_help
  fi
  if [ "$param" = "--force" -o "$param" = "-f" -o "$param" = "-a" ]; then
    SALPACK_FORCE=1
  fi
  if [ "$param" = "--watch" -o "$param" = "-w" -o "$param" = "-a" ]; then
    SALPACK_WATCH=1
  fi
  if [ "$param" = "--status" -o "$param" = "-s" ]; then
    SALPACK_STATUS=1
  fi
  if [ "$param" = "--check" -o "$param" = "-c" ]; then
    SALPACK_CHECK=1
  fi
  if [ "$param" = "--clone" -o "$param" = "clone" ]; then
    clone_parts "$2" "$3"
    clone_md "$2" "$3" "$4"
    echo "Do not forget to run 'grub2-install /dev/$tgt'"
    echo "after rebuild of root device."
  fi
  if [ "$param" = "--mirror" -o "$param" = "--add" -o "$param" = "add" ]; then
    clone_md "$2" "$3" "$4"
    exit 0
  fi
  if [ "$param" == "--write-mostly" -o "$param" = "-W" ]; then
    shift
    for disk in "$@"; do
      update_write_mostly "$disk"
    done
    exit 0
  fi
done

WATCHFILES="/proc/loadavg /proc/mdstat /proc/sys/dev/raid/speed_limit_min /proc/sys/dev/raid/speed_limit_max"
if [ -f /etc/hostname ]; then
  WATCHFILES="/etc/hostname $WATCHFILES"
fi

cat /proc/mdstat
echo_hr

if [ "$SALPACK_CHECK" ]; then
  for MD in $MDS; do
    echo check > /sys/block/$MD/md/sync_action
  done
  watch 'cat /proc/mdstat; grep ^ /sys/block/md*/md/mismatch_cnt'
  exit
fi

SLEEP=0
for MD in $MDS; do
  if [ -d $SYSDIR/$MD/md ]; then
    degraded=`getparam degraded`
    devs=`getdevs`
    for dev in $devs; do
      if hasstate $dev faulty; then
        if hasstate $dev write_mostly; then
          WRITE_MOSTLY=" -W"
        else
          WRITE_MOSTLY=""
        fi
        CMD="mdadm /dev/$MD --remove /dev/$dev -a$WRITE_MOSTLY /dev/$dev"
        echo $CMD
        if [ "$SALPACK_FORCE" ]; then
          $CMD
          SLEEP=3
        fi
      fi
    done
  else
    echo "WARNING: $MD is not an mdadm device!"
  fi
done

if [ "$SALPACK_STATUS" ]; then
  sleep 1s
  uptime
  cat $WATCHFILES
fi

if [ "$SALPACK_WATCH" ]; then
  echo "Watch will start in 3 seconds ..."
  sleep $SLEEP
  watch "md_status $WATCHFILES"
else
  echo "watch cat $WATCHFILES"
fi