#!/usr/bin/env bash
# ============================================================================
# Worker scaling as code (Modul 9 — "Număr de replici per worker configurabil").
#
# Three modes:
#   status                 — current replica count + live queue backlog per worker
#   set <service> <n>      — scale one worker to N replicas
#   auto [--apply]         — metric-driven: read didi_queue_depth from Prometheus
#                            and compute desired replicas per policy below.
#                            Dry-run by default; --apply performs the scaling.
#
# Policy: desired = clamp( ceil(backlog / TARGET_PER_REPLICA), MIN, MAX ).
# Backpressure signal is the Prometheus gauge didi_queue_depth (fed by the
# queue-depth poller in agent-v3). No external autoscaler needed — this is the
# "echivalent" orchestration the caiet allows (E3: "orchestrare/deploy … sau
# echivalent").
#
# Usage:
#   ./scale-workers.sh status
#   ./scale-workers.sh set worker-claims 5
#   ./scale-workers.sh auto            # show plan
#   ./scale-workers.sh auto --apply    # apply plan
# ============================================================================
set -uo pipefail

cd "$(dirname "$0")"

PROM_URL="${PROM_URL:-http://localhost:9090}"
COMPOSE="docker compose"

# service | prometheus component label | MIN | MAX | TARGET messages per replica
POLICY=$(cat <<'EOF'
worker-techniques      techniques       2 8  25
worker-ai-tampered     ai_tampered      2 8  25
worker-claims          claims           3 12 20
worker-domain          domain           2 6  30
worker-media-preprocess media_preprocess 2 6 5
EOF
)

current_replicas() {
  # count running containers for a compose service
  $COMPOSE ps --status running "$1" 2>/dev/null | tail -n +2 | grep -c . 2>/dev/null || echo 0
}

queue_backlog() {
  # sum(didi_queue_depth) for a component across all plan tiers
  local comp="$1"
  curl -s --get "$PROM_URL/api/v1/query" \
    --data-urlencode "query=sum(didi_queue_depth{component=\"$comp\"})" 2>/dev/null \
    | python3 -c 'import sys,json
try:
    r=json.load(sys.stdin)["data"]["result"]
    print(int(float(r[0]["value"][1])) if r else 0)
except Exception:
    print(0)'
}

clamp() { local v=$1 lo=$2 hi=$3; (( v<lo )) && v=$lo; (( v>hi )) && v=$hi; echo "$v"; }

cmd_status() {
  printf "%-24s %-10s %-10s\n" "SERVICE" "REPLICAS" "BACKLOG"
  while read -r svc comp _min _max _tgt; do
    [ -z "$svc" ] && continue
    printf "%-24s %-10s %-10s\n" "$svc" "$(current_replicas "$svc")" "$(queue_backlog "$comp")"
  done <<< "$POLICY"
}

cmd_set() {
  local svc="$1" n="$2"
  echo "Scaling $svc → $n replicas..."
  $COMPOSE up -d --no-recreate --scale "$svc=$n" "$svc"
}

cmd_auto() {
  local apply="${1:-}"
  printf "%-24s %-8s %-8s %-8s %s\n" "SERVICE" "BACKLOG" "CURRENT" "DESIRED" "ACTION"
  while read -r svc comp min max tgt; do
    [ -z "$svc" ] && continue
    local backlog cur desired
    backlog=$(queue_backlog "$comp")
    cur=$(current_replicas "$svc")
    # ceil(backlog / tgt), then clamp to [min,max]
    desired=$(( (backlog + tgt - 1) / tgt ))
    (( desired < 1 )) && desired=1
    desired=$(clamp "$desired" "$min" "$max")
    local action="keep"
    (( desired > cur )) && action="scale-up"
    (( desired < cur )) && action="scale-down"
    printf "%-24s %-8s %-8s %-8s %s\n" "$svc" "$backlog" "$cur" "$desired" "$action"
    if [ "$apply" = "--apply" ] && [ "$action" != "keep" ]; then
      $COMPOSE up -d --no-recreate --scale "$svc=$desired" "$svc" >/dev/null 2>&1 \
        && echo "   applied: $svc → $desired" || echo "   FAILED: $svc"
    fi
  done <<< "$POLICY"
  [ "$apply" != "--apply" ] && echo "(dry-run — pass --apply to scale)"
}

case "${1:-status}" in
  status) cmd_status ;;
  set)    [ $# -eq 3 ] || { echo "usage: $0 set <service> <n>"; exit 1; }; cmd_set "$2" "$3" ;;
  auto)   cmd_auto "${2:-}" ;;
  *) echo "usage: $0 {status|set <service> <n>|auto [--apply]}"; exit 1 ;;
esac