#!/usr/bin/env bash
# pod-ram-usage.sh — show RAM usage by pod, with optional namespace / sorting flags
#
# Usage:
#   ./pod-ram-usage.sh                        # all namespaces
#   ./pod-ram-usage.sh -n my-namespace        # specific namespace
#   ./pod-ram-usage.sh -s                     # sort by usage (highest first)
#   ./pod-ram-usage.sh -n kube-system -s      # combine flags
#   ./pod-ram-usage.sh -w 5                   # watch mode, refresh every 5 seconds
#   ./pod-ram-usage.sh -h                     # show help

set -euo pipefail

# ── defaults ────────────────────────────────────────────────────────────────
NAMESPACE=""
SORT=false
WATCH_INTERVAL=0   # 0 = run once

# ── helpers ─────────────────────────────────────────────────────────────────
usage() {
  grep '^#' "$0" | grep -v '#!/' | sed 's/^# \{0,2\}//'
  exit 0
}

die() { echo "ERROR: $*" >&2; exit 1; }

check_deps() {
  command -v kubectl >/dev/null 2>&1 || die "kubectl not found in PATH"
  # metrics-server check (soft): warn rather than abort
  if ! kubectl top pod --all-namespaces >/dev/null 2>&1; then
    echo "WARNING: 'kubectl top' failed — is metrics-server installed and ready?" >&2
    echo "         Install it with:" >&2
    echo "           kubectl apply -f https://github.com/kubernetes-sigs/metrics-server/releases/latest/download/components.yaml" >&2
    exit 1
  fi
}

# Convert Mi/Ki/Gi strings to MiB float for sorting
to_mib() {
  local val="$1"
  if   [[ "$val" =~ ^([0-9]+(\.[0-9]+)?)Gi$ ]]; then echo "${BASH_REMATCH[1]} * 1024"    | bc
  elif [[ "$val" =~ ^([0-9]+(\.[0-9]+)?)Mi$ ]]; then echo "${BASH_REMATCH[1]}"
  elif [[ "$val" =~ ^([0-9]+(\.[0-9]+)?)Ki$ ]]; then echo "scale=2; ${BASH_REMATCH[1]} / 1024" | bc
  elif [[ "$val" =~ ^([0-9]+(\.[0-9]+)?)m$  ]]; then echo "0"   # millicores — shouldn't appear for RAM
  else echo "0"
  fi
}

# ── parse args ───────────────────────────────────────────────────────────────
while getopts ":n:sw:h" opt; do
  case $opt in
    n) NAMESPACE="$OPTARG" ;;
    s) SORT=true ;;
    w) WATCH_INTERVAL="$OPTARG" ;;
    h) usage ;;
    :) die "Option -$OPTARG requires an argument." ;;
    \?) die "Unknown option -$OPTARG" ;;
  esac
done

# ── main logic ───────────────────────────────────────────────────────────────
run_once() {
  # build kubectl top args
  local top_args=("top" "pod")
  if [[ -n "$NAMESPACE" ]]; then
    top_args+=("-n" "$NAMESPACE")
  else
    top_args+=("--all-namespaces")
  fi
  top_args+=("--no-headers")

  # fetch raw data
  local raw
  raw=$(kubectl "${top_args[@]}" 2>/dev/null) || {
    echo "ERROR: kubectl top pod failed." >&2; return 1
  }

  # also get resource requests/limits for context
  local req_args=("get" "pods")
  if [[ -n "$NAMESPACE" ]]; then
    req_args+=("-n" "$NAMESPACE")
  else
    req_args+=("--all-namespaces")
  fi
  req_args+=("-o" "custom-columns=\
NS:.metadata.namespace,\
NAME:.metadata.name,\
MEM_REQ:.spec.containers[*].resources.requests.memory,\
MEM_LIM:.spec.containers[*].resources.limits.memory" \
    "--no-headers")

  local limits
  limits=$(kubectl "${req_args[@]}" 2>/dev/null) || limits=""

  # ── header ──────────────────────────────────────────────────────────────
  local ns_label="ALL NAMESPACES"
  [[ -n "$NAMESPACE" ]] && ns_label="$NAMESPACE"
  echo ""
  echo "═══════════════════════════════════════════════════════════════════════"
  printf "  RAM Usage by Pod — namespace: %-35s\n" "$ns_label"
  printf "  %s\n" "$(date '+%Y-%m-%d %H:%M:%S')"
  echo "═══════════════════════════════════════════════════════════════════════"
  printf "%-20s  %-42s  %10s  %10s  %10s\n" \
    "NAMESPACE" "POD" "MEMORY" "REQUEST" "LIMIT"
  echo "───────────────────────────────────────────────────────────────────────"

  # ── build output rows ────────────────────────────────────────────────────
  declare -a rows
  while IFS= read -r line; do
    [[ -z "$line" ]] && continue

    # kubectl top --all-namespaces: NAMESPACE  POD  CPU  MEM
    # kubectl top -n NS:            POD  CPU  MEM
    if [[ -n "$NAMESPACE" ]]; then
      local pod cpu mem
      read -r pod cpu mem <<< "$line"
      local ns="$NAMESPACE"
    else
      local ns pod cpu mem
      read -r ns pod cpu mem <<< "$line"
    fi

    # look up request / limit from the limits table
    local req lim
    req=$(echo "$limits" | awk -v n="$ns" -v p="$pod" '$1==n && $2==p {print $3}' | head -1)
    lim=$(echo "$limits" | awk -v n="$ns" -v p="$pod" '$1==n && $2==p {print $4}' | head -1)
    [[ -z "$req" || "$req" == "<none>" ]] && req="—"
    [[ -z "$lim" || "$lim" == "<none>" ]] && lim="—"

    # for sort key convert mem to MiB numeric
    local sort_key
    sort_key=$(to_mib "$mem")

    rows+=("$(printf '%020.4f\t%-20s\t%-42s\t%10s\t%10s\t%10s' \
      "$sort_key" "$ns" "$pod" "$mem" "$req" "$lim")")
  done <<< "$raw"

  # ── sort if requested ────────────────────────────────────────────────────
  if $SORT; then
    IFS=$'\n' sorted=($(printf '%s\n' "${rows[@]}" | sort -t$'\t' -k1 -rn))
  else
    IFS=$'\n' sorted=("${rows[@]}")
  fi

  # ── print rows ───────────────────────────────────────────────────────────
  local total_mib=0
  local count=0
  for row in "${sorted[@]}"; do
    local sort_key ns pod mem req lim
    IFS=$'\t' read -r sort_key ns pod mem req lim <<< "$row"
    printf "%-20s  %-42s  %10s  %10s  %10s\n" "$ns" "$pod" "$mem" "$req" "$lim"
    total_mib=$(echo "$total_mib + $sort_key" | bc)
    (( count++ )) || true
  done

  echo "───────────────────────────────────────────────────────────────────────"
  printf "  %d pod(s)   Total usage: %.0f MiB\n" "$count" "$total_mib"
  echo "═══════════════════════════════════════════════════════════════════════"
  echo ""
}

# ── entry point ──────────────────────────────────────────────────────────────
check_deps

if [[ "$WATCH_INTERVAL" -gt 0 ]]; then
  echo "Watch mode — refreshing every ${WATCH_INTERVAL}s (Ctrl-C to stop)"
  while true; do
    clear
    run_once
    sleep "$WATCH_INTERVAL"
  done
else
  run_once
fi