# filename: diskspace.sh
#!/usr/bin/env bash
set -euo pipefail

SERVERS=(
  "controller.dgse.cloud"
  "worker-0.dgse.cloud"
  "worker-1.dgse.cloud"
)

SSH_OPTS=(
  -o BatchMode=no
  -o ConnectTimeout=8
  -o StrictHostKeyChecking=accept-new
  -o ServerAliveInterval=10
  -o ServerAliveCountMax=2
)

EXCLUDE_TYPES_REGEX='^(tmpfs|devtmpfs|overlay|squashfs|proc|sysfs|cgroup2?|pstore|rpc_pipefs|nsfs|bpf)$'
EXCLUDE_MOUNTS_REGEX='^/(dev|proc|sys|run)(/|$)'

to_mib() {
  local val num u
  val="$(echo "$1" | awk '{print toupper($0)}')"
  if [[ "$val" =~ ^([0-9]*\.?[0-9]+)([KMGTPE]?)(I?B)?$ ]]; then
    num="${BASH_REMATCH[1]}"
    u="${BASH_REMATCH[2]}"
    case "$u" in
      "")   printf "%.0f\n" "$num" ;;
      K)    awk -v n="$num" 'BEGIN{printf "%.0f\n", n/1024}' ;;
      M)    printf "%.0f\n" "$num" ;;
      G)    awk -v n="$num" 'BEGIN{printf "%.0f\n", n*1024}' ;;
      T)    awk -v n="$num" 'BEGIN{printf "%.0f\n", n*1024*1024}' ;;
      P)    awk -v n="$num" 'BEGIN{printf "%.0f\n", n*1024*1024*1024}' ;;
      E)    awk -v n="$num" 'BEGIN{printf "%.0f\n", n*1024*1024*1024*1024}' ;;
      *)    echo 0 ;;
    esac
  else
    echo 0
  fi
}

mib_pretty() {
  local mib="$1"
  if (( mib < 1024 )); then
    printf "%d MiB" "$mib"
  elif (( mib < 1024*1024 )); then
    awk -v m="$mib" 'BEGIN{printf "%.2f GiB", m/1024}'
  else
    awk -v m="$mib" 'BEGIN{printf "%.2f TiB", m/1048576}'
  fi
}

ssh_run() {
  local user="$1" pass="$2" host="$3" cmd="$4"
  if command -v sshpass >/dev/null 2>&1; then
    SSHPASS="$pass" sshpass -e ssh "${SSH_OPTS[@]}" "${user}@${host}" "$cmd"
  else
    ssh "${SSH_OPTS[@]}" "${user}@${host}" "$cmd"
  fi
}

check_host() {
  local user="$1" pass="$2" host="$3"
  local cmd='df -hPT || df -hP'
  local output
  if ! output="$(ssh_run "$user" "$pass" "$host" "$cmd" 2>&1)"; then
    echo "ERROR: Failed to connect or run df on $host" >&2
    echo "$output" >&2
    # Emit zero so the caller can still aggregate safely
    echo "__TOTAL__ $host 0 0"
    return 1
  fi

  local total_mib=0
  local avail_mib=0
  local lines=0

  while IFS= read -r line; do
    [[ "$line" =~ ^Filesystem[[:space:]] ]] && continue
    read -r fs type size used avail usep mount <<<"$line" || true
    if [[ -z "$mount" ]]; then
      read -r fs size used avail usep mount <<<"$line" || true
      type="unknown"
    fi
    [[ -z "$fs" || -z "$size" || -z "$avail" || -z "$mount" ]] && continue
    [[ "$type" =~ $EXCLUDE_TYPES_REGEX ]] && continue
    [[ "$mount" =~ $EXCLUDE_MOUNTS_REGEX ]] && continue

    local size_mib avail_mib_i
    size_mib="$(to_mib "$size")"
    avail_mib_i="$(to_mib "$avail")"
    (( size_mib <= 0 )) && continue

    total_mib=$(( total_mib + size_mib ))
    avail_mib=$(( avail_mib + avail_mib_i ))
    lines=$(( lines + 1 ))
  done <<< "$output"

  echo "Host: $host"
  if (( lines == 0 )); then
    echo "  No eligible filesystems found (might be a minimal container or all mounts excluded)."
    echo
  else
    echo "  Total capacity:    $(mib_pretty "$total_mib")"
    echo "  Remaining (free):  $(mib_pretty "$avail_mib")"
    echo
  fi
  # Emit a parsable summary line for aggregation

  overall_avail=$(( overall_avail + avail_mib ))
  echo "__TOTAL__ $host $total_mib $avail_mib"
}

main() {
  local username
  read -r -p "SSH username: " username
  echo

  local overall_total=0
  local overall_avail=0
  local ok_hosts=0

  # Call check_host for each and parse the emitted summary lines
  for host in "${SERVERS[@]}"; do
    check_host "$username" "none" "$host"
  done | while read -r tag host total avail; do
    if [[ "$tag" == "__TOTAL__" ]]; then
      overall_total=$(( overall_total + total ))
      overall_avail=$(( overall_avail + avail ))
      (( total > 0 || avail > 0 )) && ok_hosts=$(( ok_hosts + 1 ))
    else
      # Pass through the human-readable lines
      echo "$tag $host $total $avail"
    fi
  done

  echo "Summary:"
  echo "  Hosts processed: $ok_hosts/${#SERVERS[@]}"
  echo "  Combined total capacity:   $(mib_pretty "$overall_total")"
  echo "  Combined remaining (free): $(mib_pretty "$overall_avail")"
}

main "$@"