#!/usr/bin/env bash
set -eo pipefail

readonly PROVISIONER="apelocal.csi.apecloud.com"

fail() {
  printf 'ERROR: %s\n' "$*" >&2
  exit 1
}

usage() {
  cat <<'EOF'
Usage:
  apelocal.sh check [--nodes nodes.txt] [--ssh-port PORT]
  apelocal.sh install --image-registry HOST/PATH --device DEVICE \
    [--chart FILE.tgz] [--nodes nodes.txt] [--ssh-port PORT]
  apelocal.sh install --image-registry HOST/PATH --no-storage-class \
    [--chart FILE.tgz] [--nodes nodes.txt]

Check and install use every node returned by kubectl unless --nodes limits the
operation. A nodes file contains one exact node name per non-empty line; the
same name is used as the SSH target.
EOF
}

(( $# > 0 )) || { usage; exit 1; }
if [[ "$1" == "-h" || "$1" == "--help" ]]; then
  usage
  exit 0
fi
COMMAND=$1
shift

NODES_FILE=""
SSH_PORT=""
IMAGE_REGISTRY=""
DEVICE_NAME=""
CHART_FILE=""
CREATE_STORAGE_CLASS=1
readonly RELEASE_NAME="ape-local"
readonly NAMESPACE="kb-system"
readonly STORAGE_CLASS="kbe-local"
readonly NODE_LABEL="apelocal.apecloud.com/enabled"
readonly MINIMUM_CAPACITY="100Gi"
readonly ALLOCATABLE_CAPACITY_PERCENT="120"

while (( $# )); do
  case "$1" in
    --nodes)
      (( $# >= 2 )) || fail "--nodes requires a path"
      NODES_FILE=$2
      shift 2
      ;;
    --image-registry)
      (( $# >= 2 )) || fail "--image-registry requires a value"
      IMAGE_REGISTRY=$2
      shift 2
      ;;
    --ssh-port)
      (( $# >= 2 )) || fail "--ssh-port requires a value"
      SSH_PORT=$2
      shift 2
      ;;
    --device)
      (( $# >= 2 )) || fail "--device requires a disk name"
      DEVICE_NAME=$2
      shift 2
      ;;
    --no-storage-class)
      CREATE_STORAGE_CLASS=0
      shift
      ;;
    --chart)
      (( $# >= 2 )) || fail "--chart requires a path"
      CHART_FILE=$2
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      fail "unknown argument: $1"
      ;;
  esac
done

[[ "${COMMAND}" == "check" || "${COMMAND}" == "install" ]] ||
  fail "subcommand must be check or install"

if [[ "${COMMAND}" == "check" ]]; then
  [[ -z "${IMAGE_REGISTRY}" ]] ||
    fail "check does not accept --image-registry"
  [[ -z "${DEVICE_NAME}" ]] || fail "check does not accept --device"
  (( CREATE_STORAGE_CLASS == 1 )) ||
    fail "check does not accept --no-storage-class"
  [[ -z "${CHART_FILE}" ]] || fail "check does not accept --chart"
fi

[[ -z "${NODES_FILE}" || -r "${NODES_FILE}" ]] ||
  fail "--nodes must name a readable text file when set"
command -v kubectl >/dev/null || fail "missing command: kubectl"
if [[ "${COMMAND}" == "check" ]] ||
   { [[ "${COMMAND}" == "install" ]] && (( CREATE_STORAGE_CLASS == 1 )); }; then
  command -v ssh >/dev/null || fail "missing command: ssh"
fi
if [[ "${COMMAND}" == "install" ]]; then
  command -v helm >/dev/null || fail "missing command: helm"
  [[ "${IMAGE_REGISTRY}" =~ ^[A-Za-z0-9][A-Za-z0-9._:-]*(/[A-Za-z0-9._-]+)*$ ]] ||
    fail "--image-registry must not contain a scheme or credentials"
  if (( CREATE_STORAGE_CLASS == 1 )); then
    [[ "${DEVICE_NAME}" =~ ^[A-Za-z0-9_-]+$ ]] ||
      fail "install requires a simple --device disk name"
    LOCAL_DEVICE="/dev/${DEVICE_NAME}"
  else
    [[ -z "${DEVICE_NAME}" ]] ||
      fail "--device cannot be used with --no-storage-class"
    [[ -z "${SSH_PORT}" ]] ||
      fail "--ssh-port cannot be used with --no-storage-class"
  fi
  if [[ -n "${CHART_FILE}" ]]; then
    [[ "${CHART_FILE}" == *.tgz && -f "${CHART_FILE}" &&
       -r "${CHART_FILE}" ]] ||
      fail "--chart must name a readable .tgz file"
    [[ "${CHART_FILE}" == /* ]] || CHART_FILE="$(pwd)/${CHART_FILE}"
  fi
fi

if [[ -n "${SSH_PORT}" ]]; then
  [[ "${SSH_PORT}" =~ ^[1-9][0-9]{0,4}$ ]] &&
    (( 10#${SSH_PORT} <= 65535 )) ||
    fail "--ssh-port must be an integer from 1 to 65535"
fi
declare -a SSH_ARGS=(
  -o BatchMode=yes
  -o PasswordAuthentication=no
  -o KbdInteractiveAuthentication=no
  -o StrictHostKeyChecking=accept-new
)
[[ -z "${SSH_PORT}" ]] || SSH_ARGS+=(-p "${SSH_PORT}")
readonly -a SSH_ARGS

declare -a NODES=()
if [[ -n "${NODES_FILE}" ]]; then
  node_source=$(<"${NODES_FILE}")
else
  node_source=$(kubectl get nodes \
    -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}')
fi
while IFS= read -r node || [[ -n "${node}" ]]; do
  [[ -z "${node}" ]] && continue
  [[ "${node}" =~ ^[a-z0-9]([-a-z0-9.]*[a-z0-9])?$ ]] ||
    fail "invalid Kubernetes node name"
  for existing_node in "${NODES[@]}"; do
    [[ "${existing_node}" != "${node}" ]] ||
      fail "duplicate Kubernetes node name: ${node}"
  done
  NODES+=("${node}")
done <<<"${node_source}"
(( ${#NODES[@]} > 0 )) || fail "no Kubernetes nodes selected"

node_state() {
  local node=$1 ready unschedulable
  ready=$(kubectl get node "${node}" \
    -o jsonpath='{range .status.conditions[?(@.type=="Ready")]}{.status}{end}')
  unschedulable=$(kubectl get node "${node}" -o jsonpath='{.spec.unschedulable}')
  printf '%s\t%s\n' "${ready:-Unknown}" "${unschedulable:-false}"
}

if [[ "${COMMAND}" == "check" ]]; then
  for node in "${NODES[@]}"; do
    state=$(node_state "${node}")
    printf '\nNODE %s  Ready=%s  unschedulable=%s\n' \
      "${node}" "${state%%$'\t'*}" "${state#*$'\t'}"
    ssh "${SSH_ARGS[@]}" -- "${node}" \
      "lsblk -d -e 7 -o NAME,SIZE,TYPE,FSTYPE,MOUNTPOINT ||
  printf 'WARNING: lsblk summary failed\n' >&2"
  done
  exit 0
fi

for node in "${NODES[@]}"; do
  state=$(node_state "${node}")
  [[ "${state%%$'\t'*}" == "True" ]] ||
    fail "node ${node} is not Ready"
  [[ "${state#*$'\t'}" == "false" ]] ||
    fail "node ${node} is unschedulable"
  if (( CREATE_STORAGE_CLASS == 1 )); then
    ssh "${SSH_ARGS[@]}" -- "${node}" \
      "set -eu
fail_disk() { printf 'ERROR: %s\n' \"\$1\" >&2; exit 1; }
device='${LOCAL_DEVICE}'
test -b \"\${device}\" || fail_disk \"\${device} is not a block device\"
test \"\$(lsblk -dn -o TYPE -- \"\${device}\" | tr -d '[:space:]')\" = disk ||
  fail_disk \"\${device} is not a whole disk\"
test \"\$(lsblk -nrpo NAME -- \"\${device}\" | wc -l | tr -d '[:space:]')\" = 1 ||
  fail_disk \"\${device} has partitions or child devices\"
device_basename=\${device##*/}
holder_dir=\"/sys/class/block/\${device_basename}/holders\"
test -d \"\${holder_dir}\" ||
  fail_disk \"cannot inspect block holders for \${device}\"
holders=\$(find \"\${holder_dir}\" -mindepth 1 -maxdepth 1 -print -quit 2>/dev/null) ||
  fail_disk \"cannot inspect block holders for \${device}\"
test -z \"\${holders}\" ||
  fail_disk \"\${device} has active block holders\"
test -z \"\$(findmnt -rn -S \"\${device}\" || true)\" ||
  fail_disk \"\${device} is mounted\"
test -z \"\$(wipefs -n -- \"\${device}\")\" ||
  fail_disk \"\${device} has a filesystem, partition, RAID, or LVM signature\"
if command -v pvs >/dev/null 2>&1 &&
   pvs --noheadings -o pv_name 2>/dev/null | tr -d ' ' | grep -Fxq \"\${device}\"; then
  fail_disk \"\${device} is an LVM physical volume\"
fi"
  fi
done

if helm -n "${NAMESPACE}" status "${RELEASE_NAME}" >/dev/null 2>&1; then
  fail "Helm release ${NAMESPACE}/${RELEASE_NAME} already exists"
fi
if (( CREATE_STORAGE_CLASS == 1 )) &&
   kubectl get storageclass "${STORAGE_CLASS}" >/dev/null 2>&1; then
  fail "StorageClass ${STORAGE_CLASS} already exists"
fi

node_is_reviewed() {
  local candidate=$1 node
  for node in "${NODES[@]}"; do
    [[ "${node}" != "${candidate}" ]] || return 0
  done
  return 1
}

cluster_labeled_nodes=$(kubectl get nodes -l "${NODE_LABEL}=true" \
  -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}') ||
  fail "cannot list nodes selected by ${NODE_LABEL}"
while IFS= read -r node; do
  [[ -z "${node}" ]] && continue
  node_is_reviewed "${node}" ||
    fail "labeled node ${node} is outside the reviewed inventory"
done <<<"${cluster_labeled_nodes}"

declare -a NODES_TO_LABEL=()
label_key_escaped=${NODE_LABEL//./\\.}
for node in "${NODES[@]}"; do
  current_label=$(kubectl get node "${node}" \
    -o "jsonpath={.metadata.labels.${label_key_escaped}}" 2>/dev/null || true)
  [[ -z "${current_label}" || "${current_label}" == "true" ]] ||
    fail "node ${node} has a conflicting ${NODE_LABEL} label"
  [[ -n "${current_label}" ]] || NODES_TO_LABEL+=("${node}")
done

work_dir=$(mktemp -d "${TMPDIR:-/tmp}/apelocal.XXXXXX")
declare -a LABELS_ADDED=()
RELEASE_CREATED=0
install_exit() {
  local status=$? node
  trap - EXIT
  if (( status != 0 )); then
    if (( RELEASE_CREATED == 1 )); then
      if ! helm -n "${NAMESPACE}" uninstall "${RELEASE_NAME}"; then
        printf 'WARNING: failed to roll back Helm release %s/%s\n' \
          "${NAMESPACE}" "${RELEASE_NAME}" >&2
      fi
    fi
    for node in "${LABELS_ADDED[@]}"; do
      if ! kubectl label node "${node}" "${NODE_LABEL}-"; then
        printf 'WARNING: failed to roll back label %s on node %s\n' \
          "${NODE_LABEL}" "${node}" >&2
      fi
    done
  fi
  rm -rf -- "${work_dir}"
  exit "${status}"
}
trap install_exit EXIT
if [[ -n "${CHART_FILE}" ]]; then
  CHART=${CHART_FILE}
  expected_chart_name=ape-local-csi-driver
else
  helm repo update
  chart_search="${work_dir}/chart-search.txt"
  helm search repo ape-local-csi-driver >"${chart_search}"
  chart_matches=$(
    awk 'index($1, "/") && $1 ~ /ape-local-csi-driver$/ { print $1, $2 }' \
      "${chart_search}" |
      LC_ALL=C sort -u
  )
  [[ "$(printf '%s\n' "${chart_matches}" | sed '/^$/d' | wc -l | tr -d ' ')" == "1" ]] ||
    fail "helm search repo ape-local-csi-driver must find exactly one Chart"
  read -r CHART CHART_VERSION <<<"${chart_matches}"
  expected_chart_name=${CHART#*/}
fi
chart_metadata="${work_dir}/chart.yaml"
declare -a CHART_VERSION_ARGS=()
if [[ -n "${CHART_FILE}" ]]; then
  helm show chart "${CHART}" >"${chart_metadata}"
else
  CHART_VERSION_ARGS=(--version "${CHART_VERSION}")
  helm show chart "${CHART}" \
    "${CHART_VERSION_ARGS[@]}" >"${chart_metadata}"
fi
actual_chart_name=$(awk '$1 == "name:" { print $2; exit }' "${chart_metadata}")
[[ "${actual_chart_name}" == "${expected_chart_name}" ]] ||
  fail "unexpected ape-local Chart name"
actual_chart_version=$(awk '$1 == "version:" { print $2; exit }' "${chart_metadata}")
if [[ -n "${CHART_FILE}" ]]; then
  [[ -n "${actual_chart_version}" ]] ||
    fail "local ape-local Chart omits version"
  CHART_VERSION=${actual_chart_version}
else
  [[ "${actual_chart_version}" == "${CHART_VERSION}" ]] ||
    fail "unexpected ape-local Chart version"
fi
chart_values="${work_dir}/chart-values.yaml"
helm show values "${CHART}" \
  "${CHART_VERSION_ARGS[@]}" >"${chart_values}"
grep -Eq '^image:[[:space:]]*$' "${chart_values}" ||
  fail "Chart values omit image"
grep -Eq '^[[:space:]]+registry:' "${chart_values}" ||
  fail "Chart values omit image.registry"
grep -Eq '^[[:space:]]+nodeSelector:' "${chart_values}" ||
  fail "Chart values omit node.nodeSelector"

for node in "${NODES_TO_LABEL[@]}"; do
  kubectl label node "${node}" "${NODE_LABEL}=true" --overwrite=false
  LABELS_ADDED+=("${node}")
done

values_file="${work_dir}/values.yaml"
{
  printf 'image:\n  registry: %s\n' "${IMAGE_REGISTRY}"
  printf 'csi:\n  createStorageClasses: false\n'
  printf 'node:\n  nodeSelector:\n    %s: "true"\n' "${NODE_LABEL}"
} >"${values_file}"

helm install "${RELEASE_NAME}" "${CHART}" \
  "${CHART_VERSION_ARGS[@]}" --namespace "${NAMESPACE}" --create-namespace \
  -f "${values_file}" --atomic --wait --wait-for-jobs --hide-notes --timeout 10m
RELEASE_CREATED=1

daemonsets=$(kubectl -n "${NAMESPACE}" get daemonset \
  -l "app.kubernetes.io/instance=${RELEASE_NAME}" \
  -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}')
[[ "$(printf '%s\n' "${daemonsets}" | sed '/^$/d' | wc -l | tr -d ' ')" == "1" ]] ||
  fail "expected exactly one ape-local node DaemonSet"
selector=$(kubectl -n "${NAMESPACE}" get daemonset "${daemonsets}" \
  -o "jsonpath={.spec.template.spec.nodeSelector.${label_key_escaped}}")
[[ "${selector}" == "true" ]] || fail "ape-local DaemonSet nodeSelector mismatch"
images=$(kubectl -n "${NAMESPACE}" get daemonset "${daemonsets}" \
  -o jsonpath='{range .spec.template.spec.initContainers[*]}{.image}{"\n"}{end}{range .spec.template.spec.containers[*]}{.image}{"\n"}{end}')
[[ -n "${images}" ]] || fail "ape-local DaemonSet contains no images"
while IFS= read -r image; do
  [[ -z "${image}" || "${image}" == "${IMAGE_REGISTRY}/"* ]] ||
    fail "ape-local DaemonSet image is outside --image-registry"
done <<<"${images}"

if (( CREATE_STORAGE_CLASS == 1 )); then
  kubectl create -f - <<EOF
apiVersion: storage.k8s.io/v1
kind: StorageClass
metadata:
  name: ${STORAGE_CLASS}
provisioner: ${PROVISIONER}
allowVolumeExpansion: true
reclaimPolicy: Delete
volumeBindingMode: WaitForFirstConsumer
parameters:
  volumeType: "HostPath"
  definition: |
    backendType: RawDisk
    rawDiskOptions:
      diskId: data
      diskSelector:
        devNamePatterns:
        - ^${DEVICE_NAME}$
        minimumCapacity: ${MINIMUM_CAPACITY}
      diskSetup:
        fsType: xfs
        enableProjectQuota: true
        extraMkfsFlags:
        - "-m"
        - "inobtcount=0,bigtime=0"
    allocatableCapacity:
      asWithMultiplier: ${ALLOCATABLE_CAPACITY_PERCENT}
EOF

  printf 'Installed %s %s and StorageClass %s\n' \
    "${CHART}" "${CHART_VERSION}" "${STORAGE_CLASS}"
else
  printf 'Installed %s %s without creating a StorageClass\n' \
    "${CHART}" "${CHART_VERSION}"
fi