#!/usr/bin/env bash
#
# docker-stacks-update - Pull latest images and restart all running
#                        Docker Compose stacks that have updates.
#
# Usage:  docker-stacks-update
#
# Works from any directory; put it anywhere in $PATH.

set -euo pipefail

# ── Terminal colours ──────────────────────────────────────────────
if command -v tput >/dev/null 2>&1 && [ -t 1 ]; then
  BOLD=$(tput bold 2>/dev/null || true)
  RED=$(tput setaf 1 2>/dev/null || true)
  GREEN=$(tput setaf 2 2>/dev/null || true)
  YELLOW=$(tput setaf 3 2>/dev/null || true)
  CYAN=$(tput setaf 6 2>/dev/null || true)
  RESET=$(tput sgr0 2>/dev/null || true)
fi

# ── State ───────────────────────────────────────────────────────────
TARGET_STACKS=()     # user-specified stacks (empty = all)
all_names=()         # pre-filter stack names (for error messages)
updated_names=()     # stack names that were updated
updated_detail=()    # parallel: "stack: service1, service2"
unchanged_names=()   # stack names already up to date
orphaned_names=()    # stacks whose compose files are gone – removed
failed_names=()      # stack names that errored

# ── Helpers ─────────────────────────────────────────────────────────
log()     { echo -e " ${CYAN}ℹ${RESET} $*"; }
ok()      { echo -e " ${GREEN}✓${RESET} $*"; }
warn()    { echo -e " ${YELLOW}⚠${RESET} $*"; }
fail()    { echo -e " ${RED}✗${RESET} $*"; }
header()  { echo ""; echo -e " ${BOLD}${CYAN}━━━ $* ━━━${RESET}"; echo ""; }

# ── Usage ─────────────────────────────────────────────────────────────
usage() {
  cat <<EOF
Usage: $(basename "$0") [options]

Pull latest images and restart all running Docker Compose stacks
that have updates.

Options:
  -h, --help              Show this help message
  -s, --stack <name>      Target a specific stack (repeatable)

No options = update all running stacks.
EOF
  exit 0
}

# ── Parse arguments ──────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
  case "$1" in
    -h|--help) usage ;;
    -s|--stack)
      shift
      [[ $# -eq 0 ]] && { echo -e " ${RED}✗${RESET} Option --stack requires a name." >&2; exit 1; }
      TARGET_STACKS+=("$1")
      shift
      ;;
    -*)
      echo -e " ${RED}✗${RESET} Unknown option: $1" >&2; usage ;;
    *)
      echo -e " ${RED}✗${RESET} Unexpected argument: $1" >&2; usage ;;
  esac
done

# ── Pre-flight ──────────────────────────────────────────────────────
if ! command -v docker &>/dev/null; then
  echo -e "${RED}✗${RESET} Error: 'docker' not found in PATH." >&2
  exit 1
fi

header "Docker Compose Stack Updater"

# ── Discover stacks ─────────────────────────────────────────────────
log "Scanning for running Docker Compose stacks ..."

stack_json=$(docker compose ls --format json 2>/dev/null || true)

if [[ -z "$stack_json" ]]; then
  warn "No running Docker Compose stacks found."
  exit 0
fi

# Parse JSON with jq (required).
if ! command -v jq &>/dev/null; then
  echo -e "${RED}✗${RESET} Error: 'jq' is required.  Install it first." >&2
  exit 1
fi

readarray -t names  < <(echo "$stack_json" | jq -r '.[].Name')
readarray -t configs < <(echo "$stack_json" | jq -r '.[].ConfigFiles')

total=${#names[@]}
echo "  ${CYAN}●${RESET} Found ${BOLD}${total}${RESET} running stack(s):"
for name in "${names[@]}"; do
  echo "     ${CYAN}·${RESET} ${name}"
done
echo ""

# ── Spinner for long operations ─────────────────────────────────────
_spin_kill=""

_spin_start() {
  local msg="$1"
  local -a ch=('-' '\' '|' '/')
  local i=0
  (
    while true; do
      echo -ne "\r  ${CYAN}${ch[i]}${RESET} ${msg} "
      i=$(( (i+1) % 4 ))
      sleep 0.15
    done
  ) &
  _spin_kill=$!
  disown 2>/dev/null || true
}

_spin_stop() {
  [[ -z "$_spin_kill" ]] && return
  kill "$_spin_kill" 2>/dev/null || true
  wait "$_spin_kill" 2>/dev/null || true
  _spin_kill=""
  echo -ne "\r\033[K"
}

# ── Filter stacks if --stack was given ───────────────────────────────
if [[ ${#TARGET_STACKS[@]} -gt 0 ]]; then
  all_names=("${names[@]}")
  filtered_names=()
  filtered_configs=()
  for i in "${!names[@]}"; do
    for t in "${TARGET_STACKS[@]}"; do
      if [[ "${names[$i]}" == "$t" ]]; then
        filtered_names+=("${names[$i]}")
        filtered_configs+=("${configs[$i]}")
        break
      fi
    done
  done
  names=("${filtered_names[@]}")
  configs=("${filtered_configs[@]}")
  if [[ ${#names[@]} -eq 0 ]]; then
    warn "No running stacks match the requested name(s)."
    echo "  ${CYAN}→${RESET} Running stacks: ${all_names[*]}"
    exit 0
  fi
  echo "  ${CYAN}→${RESET} Targeting ${BOLD}${#names[@]}${RESET} of ${BOLD}${total}${RESET} running stack(s)"
  echo ""
fi

# ── Process each stack ──────────────────────────────────────────────
for i in "${!names[@]}"; do
  stack="${names[$i]}"
  cfg_str="${configs[$i]}"

  echo -e " ${BOLD}${CYAN}──── ${stack}${RESET}"

  # Support multi-file projects (ConfigFiles is comma-separated).
  IFS=',' read -ra cfg_files <<< "$cfg_str"
  compose_args=()
  compose_dir=""
  for f in "${cfg_files[@]}"; do
    f_trimmed=$(echo "$f" | xargs)
    compose_args+=(-f "$f_trimmed")
    # Use the directory of the first config file as working directory.
    if [[ -z "$compose_dir" ]]; then
      compose_dir=$(dirname "$f_trimmed")
    fi
  done

  # Check whether the compose files still exist (orphan detection).
  missing_files=()
  for f in "${compose_args[@]}"; do
    [[ "$f" == "-f" ]] && continue
    [[ ! -e "$f" ]] && missing_files+=("$f")
  done

  if [[ ${#missing_files[@]} -gt 0 ]]; then
    warn "Orphaned stack — compose file(s) missing:"
    for m in "${missing_files[@]}"; do
      echo "     ${m}"
    done
    _spin_start "Tearing down orphaned stack..."
    if docker compose -p "$stack" down 2>/dev/null; then
      _spin_stop; ok "${stack} — orphaned stack removed"
    else
      _spin_stop
      # Fall back to direct container / network cleanup by project label.
      log "Fallback: removing containers by project label..."
      containers=$(docker container ls -q --filter label=com.docker.compose.project="$stack") || true
      if [[ -n "$containers" ]]; then
        docker container rm -f $containers 2>/dev/null || true
      fi
      networks=$(docker network ls -q --filter label=com.docker.compose.project="$stack") || true
      if [[ -n "$networks" ]]; then
        docker network rm $networks 2>/dev/null || true
      fi
      ok "${stack} — orphaned containers removed"
    fi
    orphaned_names+=("$stack")
    echo ""; continue
  fi

  # Pull with spinner (capture output for parsing later).
  _spin_start "Pulling latest images..."
  pull_out=$(docker compose "${compose_args[@]}" pull 2>&1) || {
    _spin_stop
    fail "Image pull failed for ${stack}"
    echo "$pull_out" | head -5 | while IFS= read -r line; do echo "  ${line}"; done
    failed_names+=("$stack")
    echo ""; continue
  }
  _spin_stop

  # Recreate containers with spinner.
  _spin_start "Recreating containers..."
  docker compose "${compose_args[@]}" up -d 2>&1 || {
    _spin_stop
    fail "Container recreation failed for ${stack}"
    failed_names+=("$stack"); echo ""; continue
  }
  _spin_stop

  # Determine which services (if any) received new images by parsing the
  # pull output.  Only the pull output is authoritative — up -d restarts
  # are idempotent and tell us nothing about whether images changed.
  #
  # The only reliable indicator of an actual image update is the
  # "Downloaded newer image for <service>" string.  The "Pulled" status
  # appears even when the image is already up to date (Docker shows
  # "Pulled" for the manifest check), so we ignore it.
  updated_svc=()

  while IFS= read -r line; do
    if [[ "$line" =~ Downloaded\ newer\ image\ for\ ([^:/\ \"]+) ]]; then
      updated_svc+=("${BASH_REMATCH[1]}")
    fi
  done <<< "$pull_out"

  # Deduplicate while preserving order.
  seen=""
  deduped=()
  for svc in "${updated_svc[@]}"; do
    if [[ ! " $seen " =~ [[:space:]]${svc}[[:space:]] ]]; then
      deduped+=("$svc")
      seen+=" $svc "
    fi
  done

  if [[ ${#deduped[@]} -gt 0 ]]; then
    ok "${stack} — updated: ${BOLD}${deduped[*]}${RESET}"
    updated_names+=("$stack")
    updated_detail+=("${stack}: ${deduped[*]}")
  else
    log "${stack} — already up to date"
    unchanged_names+=("$stack")
  fi

  echo ""
done

# ── Summary ─────────────────────────────────────────────────────────
header "Summary"

total_found=$(( ${#updated_names[@]} + ${#unchanged_names[@]} + ${#orphaned_names[@]} + ${#failed_names[@]} ))
echo "  ${CYAN}●${RESET} Stacks processed: ${BOLD}${total_found}${RESET}"
echo ""

if [[ ${#updated_names[@]} -gt 0 ]]; then
  echo -e " ${GREEN}✓${RESET} ${BOLD}Updated (${#updated_names[@]})${RESET}"
  for d in "${updated_detail[@]}"; do
    echo "     ${d}"
  done
  echo ""
fi

if [[ ${#unchanged_names[@]} -gt 0 ]]; then
  echo -e " ${CYAN}−${RESET} ${BOLD}Already up to date (${#unchanged_names[@]})${RESET}"
  for s in "${unchanged_names[@]}"; do
    echo "     ${CYAN}·${RESET} ${s}"
  done
  echo ""
fi

if [[ ${#orphaned_names[@]} -gt 0 ]]; then
  echo -e " ${YELLOW}⚠${RESET} ${BOLD}Removed (orphaned — ${#orphaned_names[@]})${RESET}"
  for s in "${orphaned_names[@]}"; do
    echo "     ${YELLOW}·${RESET} ${s}"
  done
  echo ""
fi

if [[ ${#failed_names[@]} -gt 0 ]]; then
  echo -e " ${RED}✗${RESET} ${BOLD}Failed (${#failed_names[@]})${RESET}"
  for s in "${failed_names[@]}"; do
    echo "     ${RED}·${RESET} ${s}"
  done
  echo ""
fi