#!/usr/bin/env bash
# brain-status — Print the inference backend state for any box.
#
# Works on:
#   - a GPU host (vLLM + OpenAI-compatible gateway + GPU)
#   - any box running Ollama
#   - any box with an OpenAI-compatible endpoint
#
# Usage:
#   ./brain-status              # human-readable report
#   ./brain-status --json       # machine-readable JSON
#   OPENAI_BASE_URL=http://<host>:8000/v1 ./brain-status   # probe a remote gateway

set -euo pipefail

# --- Config -------------------------------------------------------------------

BRAIN_ENV="${BRAIN_ENV:-/etc/agentic/inference.env}"
# shellcheck source=/dev/null
[[ -f "$BRAIN_ENV" ]] && source "$BRAIN_ENV"

OLLAMA_URL="${OLLAMA_URL:-http://127.0.0.1:11434}"
OPENAI_BASE_URL="${OPENAI_BASE_URL:-}"
OPENAI_API_KEY="${OPENAI_API_KEY:-dummy}"
INFERENCE_API="${INFERENCE_API:-}"
JSON_MODE=false
[[ "${1:-}" == "--json" ]] && JSON_MODE=true

# --- Helpers ------------------------------------------------------------------

_curl() { curl -s --max-time 5 "$@" 2>/dev/null; }

_heading() {
  $JSON_MODE && return
  printf "\n## %s\n" "$1"
}

_bool() { [[ "$1" == "true" ]] && echo "True" || echo "False"; }

# Echo a JSON array of {"id", "serving"} for an OpenAI-compatible gateway.
# Liveness is ground-truth, not just "advertised": prefer LiteLLM's /health
# (healthy/unhealthy endpoints, one call); otherwise fall back to a 1-token
# completion probe per model. serving is true/false, or null if unknown.
_gateway_status_json() {
  python3 - "${1%/}" "${2:-dummy}" <<'PY'
import sys, json, urllib.request
base, key = sys.argv[1], sys.argv[2]

def get(url, timeout=6, data=None):
    hdr = {"Authorization": f"Bearer {key}"}
    if data is not None:
        hdr["Content-Type"] = "application/json"
    req = urllib.request.Request(url, data=data, headers=hdr)
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)

models = []
try:
    models = [m["id"] for m in get(base + "/models").get("data", [])]
except Exception:
    pass

serving, used_health = {}, False
health_base = base[:-3] if base.endswith("/v1") else base
try:
    h = get(health_base + "/health", timeout=10)
    if isinstance(h, dict) and ("healthy_endpoints" in h or "unhealthy_endpoints" in h):
        used_health = True
        def names(lst):
            out = set()
            for e in lst or []:
                m = e.get("model", "")
                out.add(m)
                if "/" in m:
                    out.add(m.split("/", 1)[1])
            return out
        healthy, unhealthy = names(h.get("healthy_endpoints")), names(h.get("unhealthy_endpoints"))
        for mid in models:
            serving[mid] = True if mid in healthy else (False if mid in unhealthy else None)
except Exception:
    pass

if not used_health:
    for mid in models:
        try:
            body = json.dumps({"model": mid, "messages": [{"role": "user", "content": "hi"}],
                               "max_tokens": 1}).encode()
            serving[mid] = "choices" in get(base + "/chat/completions", timeout=8, data=body)
        except Exception:
            serving[mid] = False

print(json.dumps([{"id": mid, "serving": serving.get(mid)} for mid in models]))
PY
}

# --- Detect what's available --------------------------------------------------

HAS_GPU=false
HAS_OLLAMA=false
HAS_GATEWAY=false

command -v nvidia-smi &>/dev/null && nvidia-smi --query-gpu=name --format=csv,noheader &>/dev/null 2>&1 && HAS_GPU=true
_curl "${OLLAMA_URL}/api/tags" | grep -q '"models"' 2>/dev/null && HAS_OLLAMA=true

if [[ -n "$OPENAI_BASE_URL" ]]; then
  _curl "${OPENAI_BASE_URL}/models" | grep -q '"data"' 2>/dev/null && HAS_GATEWAY=true
else
  # probe common local gateway ports
  for port in 8000 4000; do
    if _curl "http://127.0.0.1:${port}/v1/models" | grep -q '"data"' 2>/dev/null; then
      OPENAI_BASE_URL="http://127.0.0.1:${port}/v1"
      HAS_GATEWAY=true
      break
    fi
  done
fi

# --- JSON output mode ---------------------------------------------------------

if $JSON_MODE; then
  gpu_json="null"
  if $HAS_GPU; then
    gpu_json=$(nvidia-smi --query-gpu=name,memory.used,memory.total,utilization.gpu,temperature.gpu \
      --format=csv,noheader,nounits | awk -F', ' '{printf "{\"name\":\"%s\",\"used_mib\":%s,\"total_mib\":%s,\"util_pct\":%s,\"temp_c\":%s}", $1,$2,$3,$4,$5}')
  fi

  ollama_json="[]"
  if $HAS_OLLAMA; then
    ollama_json=$(_curl "${OLLAMA_URL}/api/tags" | python3 -c "
import sys,json
d=json.load(sys.stdin)
print(json.dumps([{'name':m['name'],'size_gb':round(m.get('size',0)/1e9,2)} for m in d.get('models',[])]))
" 2>/dev/null || echo "[]")
  fi

  gateway_json="[]"
  if $HAS_GATEWAY; then
    gateway_json=$(_gateway_status_json "$OPENAI_BASE_URL" "$OPENAI_API_KEY" 2>/dev/null || echo "[]")
  fi

  python3 -c "
import json,sys
gpu = json.loads('$gpu_json') if '$gpu_json' != 'null' else None
print(json.dumps({
  'gpu': gpu,
  'ollama': {'available': $(_bool $HAS_OLLAMA), 'url': '$OLLAMA_URL', 'models': json.loads('$ollama_json')},
  'gateway': {'available': $(_bool $HAS_GATEWAY), 'url': '${OPENAI_BASE_URL:-}', 'models': json.loads('$gateway_json')},
  'inference_api': '${INFERENCE_API:-auto}',
}, indent=2))
"
  exit 0
fi

# --- Human-readable output ----------------------------------------------------

echo "================= INFERENCE BACKEND REPORT ================="

# GPU
if $HAS_GPU; then
  _heading "GPU"
  nvidia-smi --query-gpu=name,memory.used,memory.total,utilization.gpu,temperature.gpu \
    --format=csv,noheader | sed "s/^/   /"

  # GPU compute processes — classified as vLLM, Ollama, or other.
  gpu_pids=$(nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader,nounits 2>/dev/null || true)
  if [[ -n "$gpu_pids" ]]; then
    _heading "GPU compute processes"
    seen_ollama=false
    while IFS=, read -r pid mem; do
      [[ -z "$pid" ]] && continue
      pid=$(echo "$pid" | tr -d " "); mem=$(echo "$mem" | tr -d " ")
      cmd=$(tr "\0" " " < "/proc/$pid/cmdline" 2>/dev/null || true)
      # vLLM workers spawn via multiprocessing with empty cmdlines; climb to the
      # API-server parent that carries --port/--model.
      if ! echo "$cmd" | grep -qE -- '--port|--model'; then
        ppid=$(awk '/^PPid:/{print $2}' "/proc/$pid/status" 2>/dev/null || true)
        if [[ -n "$ppid" ]]; then
          parent_cmd=$(tr "\0" " " < "/proc/$ppid/cmdline" 2>/dev/null || true)
          [[ -n "$parent_cmd" ]] && cmd="$parent_cmd"
        fi
      fi

      if echo "$cmd" | grep -qiE 'ollama|/blobs/sha256'; then
        # Ollama runner — model name comes from `ollama ps`, not the cmdline.
        printf "   [ollama] runner pid=%-7s | %6s MiB\n" "$pid" "$mem"
        seen_ollama=true
      elif echo "$cmd" | grep -qE -- '--port|vllm'; then
        port=$(echo "$cmd"   | grep -oP '(?<=--port )\d+' || true)
        model=$(echo "$cmd"  | grep -oP '(?<=--model )\S+' || true)
        served=$(echo "$cmd" | grep -oP '(?<=--served-model-name ).*?(?= --)' || true)
        util=$(echo "$cmd"   | grep -oP '(?<=--gpu-memory-utilization )\S+' || true)
        printf "   [vllm]   :%-5s %-28s -> %-22s | %6s MiB | budget=%s\n" \
          "${port:-?}" "${model:-?}" "${served:-?}" "$mem" "${util:-?}"
      else
        printf "   [other]  pid=%-7s | %6s MiB\n" "$pid" "$mem"
      fi
    done <<< "$gpu_pids"

    # If an Ollama runner is on the GPU, show what model is actually loaded.
    if $seen_ollama && command -v ollama &>/dev/null; then
      ps_out=$(ollama ps 2>/dev/null | tail -n +2 || true)
      if [[ -n "$ps_out" ]]; then
        echo "   --- ollama loaded (ollama ps) ---"
        echo "$ps_out" | while read -r line; do
          [[ -n "$line" ]] && printf "   %s\n" "$line"
        done
      fi
    fi
  fi
else
  _heading "GPU"
  echo "   (no GPU detected)"
fi

# Ollama
_heading "Ollama (${OLLAMA_URL})"
if $HAS_OLLAMA; then
  _curl "${OLLAMA_URL}/api/tags" | python3 -c "
import sys,json
d=json.load(sys.stdin)
for m in d.get('models',[]):
    sz=m.get('size',0)/1e9
    print(f\"   - {m['name']:30s} {sz:.1f} GB\")
if not d.get('models'):
    print('   (no models pulled)')
" 2>/dev/null
else
  echo "   (not reachable)"
fi

# Gateway — advertised models annotated with real backend liveness.
_heading "Gateway (${OPENAI_BASE_URL:-not configured})"
if $HAS_GATEWAY; then
  _gateway_status_json "$OPENAI_BASE_URL" "$OPENAI_API_KEY" | python3 -c "
import sys,json
arr=json.load(sys.stdin)
if not arr:
    print('   (no models advertised)')
for m in arr:
    s=m['serving']
    tag='[up]  ' if s is True else ('[DOWN]' if s is False else '[?]   ')
    print(f\"   {tag} {m['id']}\")
if any(m['serving'] is False for m in arr):
    print('   WARNING: gateway advertises models whose backend is DOWN')
" 2>/dev/null
else
  echo "   (not reachable)"
fi

# Active backend
_heading "Active backend"
if [[ -n "${INFERENCE_API:-}" ]]; then
  echo "   INFERENCE_API=${INFERENCE_API}"
else
  if $HAS_GATEWAY; then
    echo "   auto-detected: openai (gateway available)"
  elif $HAS_OLLAMA; then
    echo "   auto-detected: ollama"
  else
    echo "   (no backend detected)"
  fi
fi

echo
