#!/bin/bash
set -euo pipefail

MEM_WARN_MIB=1024
MEM_CRIT_MIB=4096
SWAP_WARN_MIB=2048
SWAP_CRIT_MIB=8192
CPU_WARN=1.0
CPU_CRIT=2.0

W_NAME_VAL=40

RED="\033[0;31m"
YELLOW="\033[0;33m"
GREEN="\033[0;32m"
BLUE="\033[0;34m"
NC="\033[0m"

to_mib() { echo $(( (${1:-0}) / 1024 / 1024 )); }

human_bytes() {
  local b="${1:-0}"
  if [ -z "$b" ] || [ "$b" = "0" ]; then echo "unlimited"; return; fi
  local mib=$((b / 1024 / 1024))
  if [ "$mib" -ge 1024 ]; then echo "$((mib/1024))GiB"; else echo "${mib}MiB"; fi
}

pad_name() {
  # fixed width name (NO COLORS)
  printf "%-${W_NAME_VAL}.${W_NAME_VAL}s" "$1"
}

declare -A ST_CPU
declare -A ST_MEM
while IFS=$'\t' read -r name cpu mem; do
  name="$(echo "$name" | xargs)"
  ST_CPU["$name"]="$(echo "$cpu" | xargs)"
  ST_MEM["$name"]="$(echo "$mem" | xargs)"
done < <(docker stats --no-stream --format '{{.Name}}\t{{.CPUPerc}}\t{{.MemUsage}}' 2>/dev/null || true)

tmp="$(mktemp)"
trap 'rm -f "$tmp"' EXIT

# status|name|sev|cpu_lim|mem_lim|swap_lim|cpu_used|mem_used|oom|cid
for c in $(docker ps -aq); do
  cid="$(docker inspect -f '{{.Id}}' "$c" | cut -c1-12)"
  name="$(docker inspect -f '{{.Name}}' "$c" | sed 's#^/##')"
  status="$(docker inspect -f '{{.State.Status}}' "$c")"
  oom="$(docker inspect -f '{{.State.OOMKilled}}' "$c")"

  mem_b="$(docker inspect -f '{{.HostConfig.Memory}}' "$c")"
  swap_b="$(docker inspect -f '{{.HostConfig.MemorySwap}}' "$c")"
  mem_mib="$(to_mib "$mem_b")"
  swap_mib="$(to_mib "$swap_b")"

  mem_lim="$(human_bytes "$mem_b")"
  swap_lim="$(human_bytes "$swap_b")"

  nanocpus="$(docker inspect -f '{{.HostConfig.NanoCpus}}' "$c")"
  if [ -z "$nanocpus" ] || [ "$nanocpus" = "0" ]; then
    cpu_lim="unlimited"
  else
    cpu_lim="$(awk -v n="$nanocpus" 'BEGIN { printf "%.2f", n/1000000000 }')"
  fi

  cpu_used="${ST_CPU[$name]:--}"
  mem_used="${ST_MEM[$name]:--}"

  # severity for name
  sev="green"
  if [ "$oom" = "true" ]; then sev="red"; fi
  if [ "$mem_mib" -ge "$MEM_CRIT_MIB" ] || [ "$swap_mib" -ge "$SWAP_CRIT_MIB" ]; then sev="red"; fi

  if [ "$sev" != "red" ]; then
    if [ "$status" != "running" ]; then sev="yellow"; fi
    if [ "$mem_mib" -ge "$MEM_WARN_MIB" ] || [ "$swap_mib" -ge "$SWAP_WARN_MIB" ]; then sev="yellow"; fi
    if [ "$cpu_lim" != "unlimited" ]; then
      awk -v c="$cpu_lim" -v warn="$CPU_WARN" 'BEGIN{ exit !(c+0 >= warn) }' && sev="yellow"
      awk -v c="$cpu_lim" -v crit="$CPU_CRIT" 'BEGIN{ exit !(c+0 >= crit) }' && sev="red"
    fi
  fi

  echo "${status}|${name}|${sev}|${cpu_lim}|${mem_lim}|${swap_lim}|${cpu_used}|${mem_used}|${oom}|${cid}" >> "$tmp"
done

print_section() {
  local title="$1"
  local filter="$2"

  echo
if [ "$title" = "RUNNING" ]; then
  printf "${GREEN}===== %s =====${NC}\n" "$title"
else
  printf "${RED}===== %s =====${NC}\n" "$title"
fi

#  echo "===== ${title} ====="

  # Table without colors, aligned by column
  {
    echo "NAME|CPU_LIM|MEM_LIM|SWAP_LIM|CPU%|MEM_USED|STATUS|OOM|CONTAINER_ID"

    awk -F'|' "$filter {print \$0}" "$tmp" | sort -t'|' -k2,2 | \
    while IFS='|' read -r status name sev cpu_lim mem_lim swap_lim cpu_used mem_used oom cid; do
      name_fixed="$(pad_name "$name")"
      echo "${name_fixed}|${cpu_lim}|${mem_lim}|${swap_lim}|${cpu_used}|${mem_used}|${status}|${oom}|${cid}|${sev}"
    done
  } | column -t -s'|' | \
  while IFS= read -r line; do
    # colorize NAME only after alignment
    # last column is severity, remove it from output
    sev="$(echo "$line" | awk '{print $NF}')"
    line_no_sev="${line% *}"

    name_field="$(echo "$line_no_sev" | cut -d' ' -f1)"
    rest="${line_no_sev#"$name_field"}"

    if [ "$name_field" = "NAME" ]; then
      echo "$line_no_sev"
      continue
    fi

    case "$sev" in
      red)    printf "%b%s\n" "${RED}${name_field}${NC}" "$rest" ;;
      yellow) printf "%b%s\n" "${YELLOW}${name_field}${NC}" "$rest" ;;
      *)      printf "%b%s\n" "${GREEN}${name_field}${NC}" "$rest" ;;
    esac
  done
}


print_section "RUNNING" '$1=="running"'
print_section "NOT RUNNING" '$1!="running"'