#!/bin/bash

# Copyright UiPath 2021
#
# =================
# LICENSE AGREEMENT
# -----------------
#   Use of paid UiPath products and services is subject to the licensing agreement
#   executed between you and UiPath. Unless otherwise indicated by UiPath, use of free
#   UiPath products is subject to the associated licensing agreement available here:
#   https://www.uipath.com/legal/trust-and-security/legal-terms (or successor website).
#   You must not use this file separately from the product it is a part of or is associated with.

YELLOW='\033[0;33m'
RED='\033[0;31m'
NC='\033[0m'
CYAN='\033[0;36m'
GREEN='\033[0;32m'

function error() {
  echo -e "${RED}[ERROR][$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*${NC}\n" >&2
  exit 1
}

function info() {
  echo -e "${CYAN}[INFO] [$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*${NC}"
}

function warn() {
  echo -e "${YELLOW}[WARN] [$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*${NC}" >&2
}

function success() {
  echo -e "${GREEN}[SUCCESS] [$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*${NC}" >&2
}

function error_without_exit() {
  echo -e "${RED}[ERROR][$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*${NC}" >&2
}

function check_running_as_root() {
  if [ "$EUID" -ne 0 ]; then
    error "Please run script as root"
  fi
}

function display_usage() {
  echo
  echo "***************************************************************************************"
  echo
  echo "Utility to configure the disk for UiPath Automation Suite Installation."
  echo "Run this script to configure the disks on new machine or to extend the size of datadisk"
  echo
  echo "Arguments"
  echo "  -n|--node-type                  NodeType, Possible values: agent, server, nfs. Default to server."
  echo "  -i|--install-type               Installation mode, Possible values: online, offline. Default to online."
  echo "  -c|--cluster-disk-name          Device to host rancher and  kubelet. Ex: /dev/sdb."
  echo "  -e|--etcd-disk-name             Device to host etcd, Not required for agent node. Ex: /dev/sdb."
  echo "  -R|--ceph-raw-disk-name         Device to host ceph OSD, Not required for agent node. Ex: /dev/sdm"
  echo "  -l|--data-disk-name             Device to host datadisk, Not required for agent node. Ex: /dev/sdc."
  echo "  -b|--bundle-disk-name           Device to host the uipath bundle. "
  echo "                                    Only required for offline installation on 1st server node. "
  echo "  -k|--robot-package-disk-name    Device to host robot package cache folder"
  echo "                                    Only required for Automation Suite Robots dedicated node when 'packagecaching' is enabled"
  echo "  -P|--robot-package-path         Path to robot package cache folder"
  echo "                                    (defaults to '/uipath_asrobots_package_cache' if not set)" 
  echo "  -f|--complete-suite             Installing complete product suite or any of these products: "
  echo "                                    aicenter, apps, taskmining, documentunderstanding."
  echo "                                    This will configure the datadisk volume to be 2TiB instead of 512Gi."
  echo "  -p|--primary-server             Is this machine is first server machine? Applicable only for airgap install."
  echo "                                    This is the machine on which UiPath AutomationSuite bundle will be installed."
  echo "                                    Default to false."
  echo "  -x|--extend-data-disk           Extend the datadisk. Either attach new disk or resize the exiting datadisk."
  echo "  -r|--resize                     Used in conjunction of with --extend-data-disk to resize the exiting volume,"
  echo "                                    instead of adding new volume. "
  echo "  -d|--debug                      Run in debug."
  echo "  -h|--help                       Display help."
  echo
  echo "ExampleUsage:"
  echo
  echo "  $(basename "$0") --node-type server --install-type online \\"
  echo "    --cluster-disk-name /dev/sdb --etcd-disk-name /dev/sdc \\"
  echo "    --data-disk-name /dev/sdd"
  echo
  echo "  $(basename "$0") --data-disk-name /dev/sdh \\"
  echo "    --extend-data-disk"
  echo
  echo "***************************************************************************************"
  echo
}

UIPATH_VG="uipathvg"
DATADISK_VG="uipathdatavg"
ETCD_VG="uipathetcdvg"
BUNDLE_VG="uipathbundlevg"
ROBOT_PACKAGES_VG="robotpackagevg"

RANCHER_LV="rancherlv"
KUBELET_LV="kubeletlv"
INSTALLER_LV="installerlv"
DATADISK_LV="longhornlv"
ETCD_LV="etcdlv"
BUNDLE_LV="bundlelv"
ROBOT_PACKAGES_LV="robotpackagelv"


RANCHER_LV_SIZE=${RANCHER_LV_SIZE:-"190"}
KUBELET_LV_SIZE=${KUBELET_LV_SIZE:-"56"}
INSTALLER_LV_SIZE=${INSTALLER_LV_SIZE:-"10"}

FULL_SUITE=0
PRIMARY_SERVER=0
EXTEND_DATA_DISK=0
RESIZE=0
CLUSTER_DISK_SIZE=${CLUSTER_DISK_SIZE:-"$((RANCHER_LV_SIZE + KUBELET_LV_SIZE + INSTALLER_LV_SIZE))"}
FULL_SUITE_DATADISK_SIZE=${FULL_SUITE_DATADISK_SIZE:-"512"}
BASIC_SUITE_DATADISK_SIZE=${BASIC_SUITE_DATADISK_SIZE:-"512"}
ETCD_DISK_SIZE=${ETCD_DISK_SIZE:-"16"}
BUNDLE_DISK_SIZE=${BUNDLE_DISK_SIZE:-"512"}
ROBOT_PACKAGES_DISK_SIZE=${ROBOT_PACKAGES_DISK_SIZE:-"10"}
FORCE=""

# Required for ceph raw device configuration
CEPH_DEV_NAME=""
export PATH=$PATH:/var/lib/rancher/rke2/bin:/usr/local/bin
[[ -f "/var/lib/rancher/rke2/agent/kubelet.kubeconfig" ]] && export KUBECONFIG="/var/lib/rancher/rke2/agent/kubelet.kubeconfig"
[[ -f "/etc/rancher/rke2/rke2.yaml" ]] && export KUBECONFIG="/etc/rancher/rke2/rke2.yaml"
CEPH_CONFIG_FILE="/etc/udev/rules.d/99-ceph-raw-osd.rules"
SYMLINK_PATH="uipath/ceph"

function parse_long_args() {
  while (("$#")); do
    case "$1" in
    -d | --debug)
      info "Running in debug"
      shift
      set -x
      ;;

    -n | --node-type)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        NODE_TYPE="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;

    -i | --install-type)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        INSTALL_TYPE="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;

    -R | --ceph-raw-disk-name)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        CEPH_DEV_NAME="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;

    -c | --cluster-disk-name)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        CLUSTER_DEV_NAME="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;

    -e | --etcd-disk-name)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        ETCD_DEV_NAME="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;

    -l | --data-disk-name)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        DATADISK_DEV_NAME="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;

    -b | --bundle-disk-name)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        BUNDLE_DEV_NAME="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;

    -k | --robot-package-disk-name)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        ROBOT_PACKAGE_DEV_NAME="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;

    -P | --robot-package-path)
      if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
        ROBOT_PACKAGE_PATH="$2"
        shift 2
      else
        error "Value for $1 is missing"
        shift 2
      fi
      ;;


    -p | --primary-server)
      PRIMARY_SERVER=1
      shift
      ;;

    -f | --complete-suite)
      FULL_SUITE=1
      shift
      ;;

    -F | --force)
      info "Running with force=true"
      shift
      FORCE="yes"
      ;;

    -x | --extend-data-disk)
      EXTEND_DATA_DISK=1
      shift
      ;;

    -r | --resize)
      RESIZE=1
      shift
      ;;

    -h | --help)
      display_usage
      shift
      exit 0
      ;;

    -* | *) # unsupported flags
      display_usage
      error "Error: Unsupported flag $1"
      exit 1
      ;;
    esac
  done

  [[ -z $NODE_TYPE ]] && warn "-n|--node-type parameter is missing. Setting it as server" && NODE_TYPE="server"
  [[ -z $INSTALL_TYPE ]] && warn "-i|--install-type parameter is missing. Setting it as online" && INSTALL_TYPE="online"

  if [[ -n $ROBOT_PACKAGE_DEV_NAME ]]; then
     
    [[ -n $CLUSTER_DEV_NAME ]] && error "Error: argument -c|--cluster-disk-name cannot be set when configuring robot package cache disk"
    [[ -n $DATADISK_DEV_NAME ]] && error "Error: argument -l|--data-disk-name cannot be set when configuring robot package cache disk"
    [[ -n $ETCD_DEV_NAME ]] && error "Error: argument -e|--etcd-disk-name cannot be set when configuring robot package cache disk"
    [[ -n $BUNDLE_DEV_NAME ]] && error "Error: argument -b|--bundle-disk-name cannot be set when configuring robot package cache disk"
    [[ "$EXTEND_DATA_DISK" == "1" ]] && error "Error: cannot extend data disk when configuring robot package cache disk"
    [[ "${NODE_TYPE}" != "agent" ]] && error "Error: Argument -n|--node-type only possible value is agent, when configuring robot package cache"
    [[ -z $ROBOT_PACKAGE_PATH ]] && ROBOT_PACKAGE_PATH="/uipath_asrobots_package_cache"

     disk_size=$(get_size_in_GB "$ROBOT_PACKAGE_DEV_NAME")
      [[ "$disk_size" -lt "${ROBOT_PACKAGES_DISK_SIZE}" ]] &&
        error "Size of robot packages disk should be at least ${ROBOT_PACKAGES_DISK_SIZE}GiB. Current provisioned disk size is ${disk_size}GiB."

    return
  fi

  if [[ "$EXTEND_DATA_DISK" == 0 ]]; then
    if [[ $NODE_TYPE != "server" && -n $CEPH_DEV_NAME && $FORCE !=  "yes" ]]; then
      warn "Configuring Ceph OSD raw device on agent is not recommended. Provide --force argument to proceed"
      exit 0
    fi

    # Configure ceph disk only if provided
    [[ -n $CEPH_DEV_NAME ]] && return

    [[ -z $CLUSTER_DEV_NAME ]] && error "Error: Required argument -c|--cluster-disk-name missing" && is_partition_exist "$CLUSTER_DEV_NAME"

    [[ $NODE_TYPE == "server" ]] && [[ -z $ETCD_DEV_NAME ]] &&
      error "Error: Required argument -e|--etcd-disk-name missing" &&
      is_disk_exist "$ETCD_DEV_NAME"

    [[ $NODE_TYPE == "server" ]] && [[ -z $DATADISK_DEV_NAME ]] &&
      error "Error: Required argument -l|--data-disk-name missing" &&
      is_partition_exist "$DATADISK_DEV_NAME"

    [[ "${NODE_TYPE}" != "server" && "${NODE_TYPE}" != "agent" && "${NODE_TYPE}" != "nfs" ]] &&
      error "Error: Argument -n|--node-type possible values are agent, server and nfs"

    [[ "${INSTALL_TYPE}" != "online" && "${INSTALL_TYPE}" != "offline" ]] &&
      error "Error: Argument -i|--install-type possible values are online, offline"

    disk_size=$(get_size_in_GB "$CLUSTER_DEV_NAME")
    [[ "$disk_size" -lt "${CLUSTER_DISK_SIZE}" ]] &&
      error "Size of cluster disk should be atleast ${CLUSTER_DISK_SIZE}GiB. Current provisioned disk size is ${disk_size}GiB."

    if [[ "${NODE_TYPE}" == "server" ]]; then
      data_disk_size=$(get_size_in_GB "$DATADISK_DEV_NAME")
      [[ "${FULL_SUITE}" == "1" && "$data_disk_size" -lt "${FULL_SUITE_DATADISK_SIZE}" ]] &&
        error "Minimum requirement for complete suite datadisk is ${FULL_SUITE_DATADISK_SIZE}GiB. Current provision disk size is ${data_disk_size}GiB."

      [[ "${FULL_SUITE}" == "0" && "$data_disk_size" -lt "${BASIC_SUITE_DATADISK_SIZE}" ]] &&
        error "Minimum requirement for basic suite datadisk is ${BASIC_SUITE_DATADISK_SIZE}GiB. Current provision disk size is ${data_disk_size}GiB."

      etcd_disk_size=$(get_size_in_GB "$ETCD_DEV_NAME")
      [[ "$etcd_disk_size" -lt "${ETCD_DISK_SIZE}" ]] &&
        error "Size of etcd disk should be ${ETCD_DISK_SIZE}GiB. Current provision disk size is ${etcd_disk_size}GiB."

      if [[ "${ETCD_DEV_NAME}" == "$DATADISK_DEV_NAME" ]] || [[ "${ETCD_DEV_NAME}" == "$CLUSTER_DEV_NAME" ]]; then
        error "etcd cannot be co-hosted on datadisk or clusterdisk as etcd is very sensitive to IO"
      fi

      [[ "${INSTALL_TYPE}" == "offline" && "${PRIMARY_SERVER}" == "1" && -z $BUNDLE_DEV_NAME ]] &&
      error "Error: Required argument -b|--bundle-disk-name missing."

      [[ "${INSTALL_TYPE}" == "offline" && "${PRIMARY_SERVER}" == "1" && -n $BUNDLE_DEV_NAME ]] &&
        is_partition_exist "$BUNDLE_DEV_NAME"
    fi
  else
    [[ -z $DATADISK_DEV_NAME ]] && [[ $RESIZE == "0" ]] &&
      error "Error: Required argument -l|--data-disk-name missing" &&
      is_partition_exist "$DATADISK_DEV_NAME"
  fi
}

function get_size_in_GB() {
  local dev_name=$1
  local size_in_bytes
  size_in_bytes=$(lsblk "$dev_name" -b --nodeps --noheadings --paths --output size)
  local divisor=$((1024 * 1024 * 1024))
  local size_in_GB=$(((size_in_bytes + divisor - 1) / divisor)) # this will round up value
  echo ${size_in_GB}
}

function is_disk_exist() {
  local dev_name=$1
  lsblk --nodeps --noheadings --paths |
    grep -q "${dev_name}" ||
    error "Disk $dev_name does not exist"
}

function is_partition_exist() {
  local dev_name=$1
  lsblk  --noheadings --paths |
    grep -q "${dev_name}" ||
    error "Partition  $dev_name does not exist"
}

function create_lv() {
  local lv_name=$1
  local lv_size=$2
  local vg_name=$3

  dev_name="/dev/$vg_name/$lv_name"

  part_size=$(vgdisplay "$vg_name" -c | cut -d ":" -f13)

  parts=$((lv_size * 1024 * 1024 / part_size - 1))

  lvdisplay -c | cut -d ':' -f1 | grep -q "$dev_name" ||
    lvcreate -n "$lv_name" -l "$parts" "$vg_name" -y >/dev/null ||
    error "Unable to create logical volume for host partition: $dev_name"

  echo "$dev_name"
}

function create_vg() {
  local vg_name=$1
  local dev_name=$2

  vgdisplay -c | cut -d ':' -f1 | grep -q "$vg_name" ||
    vgcreate "$vg_name" "$dev_name" ||
    error "Unable to create volume group \"$vg_name\" from device \"$dev_name\""
}

function extend_vg() {
  local vg_name=$1
  local dev_name=$2

  if vgdisplay -c | cut -d ':' -f1 | grep -q "$vg_name"; then
    vgextend "$vg_name" "$dev_name"
  else
    error "Volume group \"$vg_name\" not found. Unable to extend volume from device \"$dev_name\""
  fi

}

function format_device() {
  local device_name=$1
  current_fs=$(lsblk -no FSTYPE "$device_name")
  [[ $current_fs == "xfs" ]] && info "$device_name is already formatted in xfs format, skipping"
  if [[ -z "${current_fs}" ]]; then
    mkfs.xfs -f "${device_name}" || error "Unable to format $device_name"
    info "$device_name formatted in xfs format"
  fi
}

function mount_partition() {
  local device_name=$1
  local mount_path=$2

  format_device "$device_name"

  if entry=$(grep -sE "(^|\s)$mount_path($|\s)" /proc/mounts); then
    mounted_dev_name=$(echo "$entry" | cut -d " " -f1)
    info "Partition: $mount_path is already mounted using Device: $mounted_dev_name, skipping"
  else
    mkdir -p "$mount_path"
    mount "${device_name}" "${mount_path}" || error "Unable to mount partition $mount_path under Device $device_name"
    success "Successfully mounted partition $mount_path using device $device_name"
  fi
}

function add_fstab_entry() {
  local dev_name=$1
  local mount_path=$2
  local part_name=$3

  uuid=$(blkid -o value -s UUID "${dev_name}")
  if [[ -z "${uuid}" ]]; then error "Unable to fetch UUID. Invalid DeviceName: $uuid"; fi

  grep -q "${mount_path}" /etc/fstab ||
    printf "# uipath-$part_name\nUUID=%s    %s    xfs    defaults,nofail    0    0\n" "${uuid}" "${mount_path}" >>/etc/fstab
}

function configure_cluster_disk() {
  local dev_name=$1
  local mount_path_rancher="/var/lib/rancher"
  local mount_path_kubelet="/var/lib/kubelet"
  local mount_path_installer="/opt/UiPathAutomationSuite"

  create_vg $UIPATH_VG "$dev_name"

  rancher_dev_name=$(create_lv $RANCHER_LV "${RANCHER_LV_SIZE}" $UIPATH_VG)
  mount_partition "$rancher_dev_name" $mount_path_rancher
  add_fstab_entry "$rancher_dev_name" $mount_path_rancher "rancher"
  info "$mount_path_rancher is configured"

  kubelet_dev_name=$(create_lv $KUBELET_LV "${KUBELET_LV_SIZE}" $UIPATH_VG)
  mount_partition "$kubelet_dev_name" $mount_path_kubelet
  add_fstab_entry "$kubelet_dev_name" $mount_path_kubelet "kubelet"
  info "$mount_path_kubelet is configured"

  installer_dev_name=$(create_lv $INSTALLER_LV "${INSTALLER_LV_SIZE}" $UIPATH_VG)
  mount_partition "$installer_dev_name" $mount_path_installer
  add_fstab_entry "$installer_dev_name" $mount_path_installer "installer"
  info "$mount_path_installer is configured"

  success "Cluster disk is configured successfully at dev location $dev_name"
}

function configure_data_disk() {
  local dev_name=$1
  local mount_path_datadisk="/datadisk"

  create_vg $DATADISK_VG "$dev_name"

  disk_size=$BASIC_SUITE_DATADISK_SIZE
  [[ "$FULL_SUITE" == 1 ]] && disk_size=$FULL_SUITE_DATADISK_SIZE

  datadisk_dev_name=$(create_lv $DATADISK_LV "${disk_size}" $DATADISK_VG)
  mount_partition "$datadisk_dev_name" $mount_path_datadisk
  add_fstab_entry "$datadisk_dev_name" $mount_path_datadisk "datadisk"
  info "$mount_path_datadisk is configured"

  success "Data disk is configured successfully at dev location $dev_name"
  extend_data_disk "${dev_name}"
}

function configure_etcd_disk() {
  local dev_name=$1
  local mount_path_etcd="/var/lib/rancher/rke2/server/db"

  create_vg "$ETCD_VG" "$dev_name"

  etcd_dev_name=$(create_lv $ETCD_LV "${ETCD_DISK_SIZE}" $ETCD_VG)
  mount_partition "$etcd_dev_name" $mount_path_etcd
  add_fstab_entry "$etcd_dev_name" $mount_path_etcd "etcd"
  info "$mount_path_etcd is configured"

  success "etcd disk is configured successfully at dev location $dev_name"
}

function configure_robot_package_disk() {
  local dev_name=$1
  local mount_path_robot_packages=$2

  create_vg "$ROBOT_PACKAGES_VG" "$dev_name"

  robot_packages_dev_name=$(create_lv $ROBOT_PACKAGES_LV "${ROBOT_PACKAGES_DISK_SIZE}" $ROBOT_PACKAGES_VG)
  mount_partition "$robot_packages_dev_name" "$mount_path_robot_packages"
  add_fstab_entry "$robot_packages_dev_name" "$mount_path_robot_packages" "uipath-robotpackages"
  info "$mount_path_robot_packages is configured"

  success "robot package disk is configured successfully at dev location $dev_name"
}

function configure_bundle_disk() {
  local dev_name=$1
  local mount_path_bundle="/uipath"
  create_vg "$BUNDLE_VG" "$dev_name"

  bundle_dev_name=$(create_lv $BUNDLE_LV "${BUNDLE_DISK_SIZE}" $BUNDLE_VG)
  mount_partition "$bundle_dev_name" $mount_path_bundle
  add_fstab_entry "$bundle_dev_name" $mount_path_bundle "bundle"
  info "$mount_path_bundle is configured"

  success "UiPath bundle disk is configured successfully at dev location $dev_name"
}

function get_device_json(){
  local device="$1"
  local device_json
  device_json=$(lsblk --bytes --output-all "${device}" --json)
  echo "${device_json}"
}

function add_device_to_ceph_config() {
  local device="$1"
  [[ -f "$CEPH_CONFIG_FILE" ]] || touch "$CEPH_CONFIG_FILE"
  sed -i "s;KERNEL==\"$device\".*;KERNEL==\"$device\", SYMLINK+=\"${SYMLINK_PATH}/${device}\";" $CEPH_CONFIG_FILE
  if ! grep -q "KERNEL==\"$device\"" $CEPH_CONFIG_FILE; then
    echo "KERNEL==\"${device}\", SYMLINK+=\"${SYMLINK_PATH}/${device}\"" >> "${CEPH_CONFIG_FILE}"
  fi
  ## Trigger device refresh
  udevadm control --reload
  udevadm trigger
}

function restart_ceph_operator_if_possible() {
  if kubectl -n rook-ceph get deploy  rook-ceph-operator >> /dev/null 2>&1; then
    kubectl -n rook-ceph rollout restart deploy/rook-ceph-operator && return 0
  fi
  return 1
}

function is_ceph_cluster_using_raw_devices() {
  local exit_code=1
  if kubectl -n rook-ceph get cephcluster rook-ceph -o json 2> /dev/null | jq '.spec.storage.devicePathFilter!=null' | grep -q 'true'; then
    exit_code=0
  fi
  return "${exit_code}"
}

function wait_for_ceph_osd_to_bootstrap() {
  local devices=("$@")
  local ceph_device
  local device_json
  local timeout="300"

  for ceph_device in "${devices[@]}"; do
    start_time=$(date +"%s")
    status="failure"
    while true; do
      current_time=$(date +"%s")
      diff=$((current_time-start_time))
      [[ $diff -gt $timeout ]] && break
      device_json=$(get_device_json "${ceph_device}")
      [[ "$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].fstype')" == "LVM2_member" && "$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].children[0].name')" == "ceph--"* ]] && status="success" && break
      is_osd_on_partition "${ceph_device}" && status="success" && break
      sleep 5
    done
    if [[ $status == "success" ]]; then
      info "OSD bootstrapped successfully on device: '${ceph_device}'"
    else
      error_without_exit "Timeout waiting for ceph OSD to bootstrap on device: '${ceph_device}'"
    fi
  done
}

function wait_until_ceph_operator_restarted() {
  kubectl -n rook-ceph rollout status deploy/rook-ceph-operator
}

function get_local_k8s_node_name(){
  local node_name
  node_name=$(xargs -0 < "/proc/$(pgrep kubelet)/cmdline" | tr ' ' '\n' | grep '\-\-hostname\-override' | cut -d'=' -f2)
  [[ -z "${node_name}" ]] && node_name=$(hostname)
  [[ -z "${node_name}" ]] && warn "Unable to get kubernetes node name"
  echo "${node_name}"
}


function is_osd_on_partition(){
  local part_name="$1"
  local node_name
  local pod_name
  node_name=$(get_local_k8s_node_name)
  [[ -z "${node_name}" ]] && return 1
  pod_name=$(kubectl -n rook-ceph get pod -l app=rook-ceph-osd -o name --field-selector spec.nodeName="${node_name}" -o jsonpath='{.items[0].metadata.name}' 2> /dev/null)
  [[ -z "${pod_name}" ]] && return 1
  kubectl -n rook-ceph exec "$pod_name" -c osd -- ceph-volume inventory "${part_name}" --format json | jq -r '.rejected_reasons | index("Has BlueStore device label")' | grep -q -E '^[0-9]+$'
}


function configure_ceph_raw_disk() {
  local device_list="$1"
  local device_list_json
  local ceph_device
  local device_json
  local base_device
  local rotational_status
  local fs_type
  local device_child
  local exit_code=0
  local timeout=30
  local is_operator_restart_required=false
  local size
  local is_device_child_present
  local eligible_device=()

  shopt -s nullglob
  device_list_json=$(jq -c -n --arg devices "${device_list}" '[$devices | split(",") | .[] | gsub("^\\s+|\\s+$";"")]')
  for ceph_device in $(jq -r -n --argjson devices "$device_list_json" '$devices[]'); do
    # Check if valid block device
    [[ ! -b "${ceph_device}" ]] && error_without_exit "Provided device: '${ceph_device}' is not a valid block device" && exit_code=1 && continue
    device_json=$(get_device_json "${ceph_device}")
    base_device=$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].name')
    [[ -z "${base_device}" ]] && error_without_exit "Empty base device for '${ceph_device}'" && exit_code=1 && continue

    # Check fs on device
    fs_type=$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].fstype')
    if [[ $fs_type != "null" ]]; then
      if [[ $fs_type == "LVM2_member" ]]; then
        device_child=$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].children[0].name')
        if [[ $device_child == "ceph--"* ]]; then
          info "Device: '${ceph_device}' is already bootstrapped for ceph OSD"
          add_device_to_ceph_config "${base_device}"
          continue
        else
          error_without_exit "Device: '${ceph_device}' is already part of non-ceph Logical Volume: '${device_child}'"
          exit_code=1
          continue
        fi
      else
        error_without_exit "Device: '${ceph_device}' contains filesystem: '${fs_type}' which is not supported for ceph OSD"
        exit_code=1
        continue
      fi
    else
      is_device_child_present=$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].children!=null')
      if [[ $is_device_child_present == "true" ]]; then
        error_without_exit "Device: '${ceph_device}' with a child device is not supported for ceph OSD"
        exit_code=1
        continue
      fi
    fi

    # Check device type, support only disk partition
    device_type=$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].type')
    if ! [[ $device_type == "disk" || $device_type == "part" ]]; then
      error_without_exit "Device: '${ceph_device}' is of type: '${device_type}' not supported. Only type: 'disk' or 'part' is supported"
      exit_code=1
      continue
    fi

    # Check if OSD already running on partition
    if [[ $device_type == "part" ]]; then
      if is_osd_on_partition "${ceph_device}"; then
        info "Device: '${ceph_device}' is already bootstrapped for ceph OSD"
        add_device_to_ceph_config "${base_device}"
        continue        
      fi
    fi

    if [[ -e "/dev/${SYMLINK_PATH}/${base_device}" ]]; then
      info "Device: '${ceph_device}' already has required DevLink"
      add_device_to_ceph_config "${base_device}"
    else
      # Check if device rotational
      rotational_status=$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].rota')
      [[ $rotational_status -ne 0 ]] && error_without_exit "Rotational device: '${ceph_device}' is not recommended for ceph OSD" && exit_code=1 && continue

      # Check if size > 5GB
      size=$(jq -r -n --argjson deviceJson "${device_json}" '$deviceJson.blockdevices[0].size')
      [[ $size -lt "5368709120" ]] && error_without_exit "Unsupported disk capacity: $(numfmt --to=iec-i "${size}") for device: '${ceph_device}'. Required > $(numfmt --to=iec-i 5368709120)" && exit_code=1 && continue

      add_device_to_ceph_config "${base_device}"
      info "Device: '${ceph_device}' has been added/updated in the config file"

      ## Check if device devlink is created
      if ! timeout $timeout bash -c "while ! [[ -L \"/dev/${SYMLINK_PATH}/${base_device}\" && -e \"/dev/${SYMLINK_PATH}/${base_device}\" && -b \"/dev/${SYMLINK_PATH}/${base_device}\" ]];do
        sleep 2
      done" ; then
        error_without_exit "Unable to configure DevLink at '/dev/${SYMLINK_PATH}/${base_device}' for Block Device: '${ceph_device}'"
        exit_code=1
        continue
      else
        info "Required DevLink have been configured for device: '${ceph_device}' successfully"
      fi
    fi
    eligible_device+=("${ceph_device}")
    is_operator_restart_required=true
  done

  if is_ceph_cluster_using_raw_devices && [[ $is_operator_restart_required == "true" ]]; then
    restart_ceph_operator_if_possible && wait_until_ceph_operator_restarted && wait_for_ceph_osd_to_bootstrap "${eligible_device[@]}"
  fi
  shopt -u nullglob
  [[ $exit_code -eq 0 ]] && success "Ceph disk is configured successfully at dev location $device_list"
  return $exit_code
}

function extend_data_disk() {
  local dev_name=$1
  local pvList
  local iter

  [[ "$RESIZE" == "0" ]] && extend_vg $DATADISK_VG "$dev_name"
  lv_dev_name="/dev/$DATADISK_VG/$DATADISK_LV"

  pvList=$(pvdisplay --select vg_name=${DATADISK_VG} --columns --options 'pv_name' --reportformat basic | awk 'NR!=1')
  for iter in $pvList ; do
    pvresize --yes --quiet "${iter}"
  done

  if lvextend -r -l +100%FREE "$lv_dev_name"; then
    success "Successfully extended partition /dev/$DATADISK_VG/$DATADISK_LV"
  else
    error "unable to extend the Datadisk volume: $lv_dev_name"
  fi
}

function detect_os() {
  # shellcheck disable=SC1091
  echo "$(source /etc/os-release ; echo "$ID")"
}

function check_package_installed_in_debian() {
    status="$(dpkg-query -W --showformat='${db:Status-Status}' "$1" &>/dev/null)" && return 1
    [[ ! "$status" == installed ]] && return 1
    
    return 0
}

function check_package_installed_in_rhel() {
    rpm -q "lvm2" &>/dev/null && return 0 || return 1
}

function is_lvm_exported_in_path() {
  lvm version &> /dev/null || error "lvm is not exported in environment variable"
}

function is_lvm_installed() {
  OS_ID="$(detect_os)"
  if [[ "${OS_ID}" == "rhel" ]]; then
    check_package_installed_in_rhel "lvm2" || error "lvm is not installed on machine"
  else
    check_package_installed_in_debian "lvm2" || error "lvm is not installed on machine"
  fi
}

function main() {
  if [[ -n "${CEPH_DEV_NAME}" ]]; then
    configure_ceph_raw_disk "${CEPH_DEV_NAME}" || exit 1
    exit 0
  fi

  is_lvm_installed && is_lvm_exported_in_path

  if [[ "$EXTEND_DATA_DISK" == "1" ]]; then
    extend_data_disk "$DATADISK_DEV_NAME" || exit 1
    exit 0
  fi

  [[ -n ${CLUSTER_DEV_NAME} ]] && configure_cluster_disk "$CLUSTER_DEV_NAME"

  [[ -n ${DATADISK_DEV_NAME} && $NODE_TYPE == "nfs" ]] && configure_data_disk "$DATADISK_DEV_NAME"

  [[ -n ${DATADISK_DEV_NAME} && $NODE_TYPE == "server" ]] && configure_data_disk "$DATADISK_DEV_NAME"

  [[ -n "$ETCD_DEV_NAME" && "$NODE_TYPE" == "server" ]] && configure_etcd_disk "$ETCD_DEV_NAME"

  [[ -n "$BUNDLE_DEV_NAME" && "$NODE_TYPE" == "server" && "$PRIMARY_SERVER" == "1" && "$INSTALL_TYPE" ]] && configure_bundle_disk "$BUNDLE_DEV_NAME"

  [[ -n "$ROBOT_PACKAGE_DEV_NAME" ]] && configure_robot_package_disk "$ROBOT_PACKAGE_DEV_NAME" "$ROBOT_PACKAGE_PATH"
  
}

check_running_as_root
parse_long_args "$@"
main
exit 0

