set -e

while getopts ":h" option; do
   case $option in
      h) # display Help
         echo "Usage: ./migrate_storage_class.sh [namespace. defaults to kfuse] [new storage class. defaults to kfuse-ssd] [tmp_dir. defaults to /tmp/kfuse_migrate_storage]"
         exit;;
   esac
done

namespace=$1
to_storage_class_name=$2
tmp_dir=$3

if [ -z "$namespace" ]; then
namespace="kfuse"
fi

if [ -z "$to_storage_class_name" ]; then
to_storage_class_name="kfuse-ssd"
fi

if [ -z "$tmp_dir" ]; then
tmp_dir="/tmp/kfuse_migrate_storage"
fi

echo Migrating storage class of Kfuse on namespace $namespace to storage class $to_storage_class_name. Using temp dir $tmp_dir.
read -p "Proceed? (y/n)"
echo    # (optional) move to a new line
if [[ ! $REPLY =~ ^[Yy]$ ]]
then
    echo Not Proceeding.
    exit
fi

stss="kafka kafka-zookeeper pinot-server-realtime pinot-server-offline pinot-zookeeper kfuse-configdb"

echo Creating temp dir $tmp_dir
mkdir -p $tmp_dir

for sts in $stss
do
  echo Checking sts $sts
  sts_info=`kubectl get sts $sts -n $namespace -o 'custom-columns=NAME:.metadata.name,STORAGE_CLASS:.spec.volumeClaimTemplates[].spec.storageClassName' --no-headers`
  has_kfuse_ssd=`echo $sts_info | awk '$2 ~ /kfuse-ssd(-aws$|-gcp$|$)/ { print $1 }'`
  if [ -z "$has_kfuse_ssd" ]; then
     echo Skipping sts $sts
     echo
     continue
  fi

  sts_existing_storage_class=`echo $sts_info | awk '{ print $2 }'`
  if [ "$sts_existing_storage_class" == "$to_storage_class_name" ]; then
    echo $sts is already using $to_storage_class_name. Skipping.
    echo
    continue
  fi

  echo Migrating storage class of sts $sts from $sts_existing_storage_class to $to_storage_class_name

  read -p "Proceed? (y/n)"
  echo    # (optional) move to a new line
  if [[ ! $REPLY =~ ^[Yy]$ ]]
  then
      echo Skipping sts $sts
      continue
  fi

  sts_pvs=`kubectl get pv -o 'custom-columns=NAME:.metadata.name,PVC:.spec.claimRef.name,PVC_NAMESPACE:.spec.claimRef.namespace,STORAGE_CLASS:.spec.storageClassName' | awk -v namespace=$namespace -v sts="data-$sts-[0-9]" '$3 == namespace && $2~sts { print $1 }'`
  for pv in $sts_pvs
  do
    echo Setting ReclaimPolicy of $pv to Retain
     kubectl patch pv $pv -n $namespace --patch '{"spec": {"persistentVolumeReclaimPolicy": "Retain"}}'
  done

  kubectl get sts -n $namespace $sts -o=yaml > $tmp_dir/$sts.yaml
  sed "s/storageClassName: $sts_existing_storage_class/storageClassName: $to_storage_class_name/g" $tmp_dir/$sts.yaml > $tmp_dir/${sts}_updated.yaml
  echo Deleting sts $sts
  kubectl delete sts $sts -n $namespace

  sts_pvcs=`kubectl get pvc -n $namespace -o 'custom-columns=NAME:.metadata.name' --no-headers | awk -v sts="data-$sts-[0-9]" '$1~sts { print $1 }'`
  for pvc in $sts_pvcs
  do
    echo Deleting pvc $pvc
    kubectl delete pvc -n $namespace $pvc
  done

  # Patch the pv to use kfuse-ssd storage class and clear claimRef
  for pv in $sts_pvs
  do
    echo Setting Storage Class of $pv to $to_storage_class_name
    kubectl patch pv $pv -n $namespace --patch "{\"spec\": {\"claimRef\": {\"uid\": null}, \"storageClassName\": \"$to_storage_class_name\"}}"
  done

  echo Applying updated sts $sts
  kubectl apply -n $namespace -f $tmp_dir/${sts}_updated.yaml

  echo Waiting for pvcs to get created back for $sts
  for pvc in $sts_pvcs
  do
    while true
    do
      pvc_status=`kubectl get pvc $pvc -n $namespace -o 'custom-columns=STATUS:.status.phase' --no-headers`
      if [ $pvc_status == "Bound" ]; then
        echo $pvc is Bound
        break
      fi
      echo $pvc is not yet Bounded
    done
  done

  for pv in $sts_pvs
  do
    echo Setting ReclaimPolicy of $pv to Delete
     kubectl patch pv $pv -n $namespace --patch '{"spec": {"persistentVolumeReclaimPolicy": "Delete"}}'
  done

  echo Finished migrating storage class of sts $sts to $to_storage_class_name
  echo
done

echo Deleting tmp dir $tmp_dir
rm -rf $tmp_dir

echo Pre-upgrade step is done. Proceed with kfuse upgrade.