#!/usr/bin/env bash
#
# Deploy the GPU serving stack on worker-1:
#   NVIDIA device plugin -> KServe (RawDeployment) -> 3 engines
#     GPU0,1  vLLM v0.11.0  deepseek-r1-distill-qwen-32b-awq  (2 replicas)
#     GPU2    Ollama        ornith:35b
#     GPU3    vLLM v0.11.0  Qwen2.5-Math-PRM-7B  (reward model)
#
# Two modes:
#   argocd (default) -- syncs the ArgoCD Applications from
#                       k8s/argocd/apps/70-gpu-serving.yaml. Git is the source
#                       of truth; this only sequences the syncs.
#   manual  (--manual) -- bootstraps directly with helm + kubectl, for first
#                       bring-up before the Applications are committed. Applies
#                       the SAME manifests from k8s/apps/llm-serving/, so ArgoCD
#                       adopts them cleanly later (matches the repo's existing
#                       k8s/bootstrap/phaseN-* pattern).
#
# Usage:
#   ./scripts/deploy-gpu-serving.sh --manual              # full manual bootstrap
#   ./scripts/deploy-gpu-serving.sh --manual gpu-plugin   # one stage
#   ./scripts/deploy-gpu-serving.sh --manual --dry-run    # show, don't run
#   ./scripts/deploy-gpu-serving.sh                       # via ArgoCD
#
set -euo pipefail

NODE=worker-1
NS=llm-serving
KSERVE_NS=kserve
GPU_NS=gpu-system
KSERVE_VER=v0.15.2       # v-PREFIXED; the unprefixed tag 404s
NVDP_VER=0.19.3
REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
MODE=argocd
DRY=false

log()  { printf '\n\033[1m==> %s\033[0m\n' "$*"; }
info() { printf '    %s\n' "$*"; }
die()  { printf '\033[31mERROR: %s\033[0m\n' "$*" >&2; exit 1; }
run()  { if $DRY; then info "[dry-run] $*"; else "$@"; fi; }
need() { command -v "$1" >/dev/null 2>&1 || die "missing required tool: $1"; }

sync_app() {
  local app=$1
  log "argocd sync: $app"
  run argocd app sync "$app" --timeout 600
  run argocd app wait "$app" --health --timeout 600
}

# ---------------------------------------------------------------- stages

preflight() {
  log "preflight"
  need kubectl
  [[ $MODE == argocd ]] && need argocd
  [[ $MODE == manual ]] && need helm

  kubectl get node "$NODE" >/dev/null 2>&1 || die "node $NODE not found"
  local ready
  ready=$(kubectl get node "$NODE" -o jsonpath='{.status.conditions[?(@.type=="Ready")].status}')
  [[ $ready == True ]] || die "node $NODE is not Ready"
  info "node $NODE: Ready"

  # Driver ships in the Talos installer image as a system extension, not from
  # the cluster. Missing label => node was built from the wrong schematic and
  # no amount of device-plugin will help.
  kubectl get node "$NODE" -o jsonpath='{.metadata.labels}' | grep -q 'nonfree-kmod-nvidia' \
    || die "NVIDIA driver extension label absent — worker-1 not on the GPU schematic"
  info "NVIDIA driver extension: present"

  kubectl get node "$NODE" -o jsonpath='{.metadata.labels.nvidia\.com/gpu}' 2>/dev/null | grep -q true \
    || die "node label nvidia.com/gpu=true missing — device plugin nodeSelector will not match"
  info "node label nvidia.com/gpu=true: present"

  [[ $(kubectl get node "$NODE" -o jsonpath='{.spec.unschedulable}') == true ]] \
    && info "node is cordoned (expected; uncordoned before engines deploy)"
  info "mode: $MODE"
}

runtimeclass() {
  # Must exist before ANY pod with runtimeClassName: nvidia is admitted —
  # including the device plugin itself. The Talos extension registers the
  # containerd handler but not this object.
  if [[ $MODE == manual ]]; then
    log "apply RuntimeClass nvidia"
    run kubectl apply -k "$REPO_ROOT/k8s/apps/gpu-runtimeclass"
  else
    sync_app gpu-runtimeclass
  fi
  $DRY || kubectl get runtimeclass nvidia >/dev/null 2>&1 \
    && info "RuntimeClass nvidia: present"
}

gpu_plugin() {
  runtimeclass
  if [[ $MODE == manual ]]; then
    log "helm install: nvidia-device-plugin $NVDP_VER"
    # Driver + container toolkit come from Talos system extensions. This chart
    # only advertises the GPUs to the kubelet — it must not install drivers.
    run helm upgrade --install nvidia-device-plugin nvidia-device-plugin \
      --repo https://nvidia.github.io/k8s-device-plugin \
      --version "$NVDP_VER" \
      -n "$GPU_NS" --create-namespace \
      --set runtimeClassName=nvidia \
      `# --set-string, NOT --set: nodeSelector is map[string]string in the` \
      `# PodSpec schema, and plain --set coerces "true" to a YAML boolean,` \
      `# which the API server rejects as a type violation.` \
      --set-string nodeSelector."nvidia\.com/gpu"=true \
      `# Drop the chart's default nodeAffinity. It requires one of three` \
      `# Node-Feature-Discovery labels (feature.node.kubernetes.io/pci-10de.present,` \
      `# .../cpu-model.vendor_id=NVIDIA, or nvidia.com/gpu.present). NFD is not` \
      `# installed here and Talos sets nvidia.com/gpu (no ".present" suffix), so` \
      `# the affinity matches zero nodes and the DaemonSet silently reports` \
      `# desiredNumberScheduled=0. The nodeSelector above is our constraint.` \
      `# Deliberately NO --wait: it blocks on DaemonSet readiness, and any` \
      `# admission failure then times helm out and wedges the release in` \
      `# pending-upgrade, blocking all later upgrades. The allocatable-GPU` \
      `# poll below is the real readiness signal.` \
      --set affinity=null
  else
    sync_app nvidia-device-plugin
  fi

  log "waiting for nvidia.com/gpu to register on $NODE"
  $DRY && { info "[dry-run] would verify allocatable == 4"; return 0; }
  local n=""
  for _ in $(seq 1 60); do
    n=$(kubectl get node "$NODE" -o jsonpath='{.status.allocatable.nvidia\.com/gpu}' 2>/dev/null || true)
    [[ -n $n && $n != 0 ]] && break
    sleep 5
  done
  [[ $n == 4 ]] || die "expected 4 allocatable GPUs, got '${n:-none}'"
  info "allocatable nvidia.com/gpu: $n"
}

kserve() {
  if [[ $MODE == manual ]]; then
    log "helm install: kserve-crd $KSERVE_VER"
    # No --wait on helm (a timeout leaves the release wedged in pending-upgrade,
    # blocking every later upgrade). kubectl wait below is the readiness signal
    # and does not touch helm release state.
    run helm upgrade --install kserve-crd "oci://ghcr.io/kserve/charts/kserve-crd" \
      --version "$KSERVE_VER" -n "$KSERVE_NS" --create-namespace

    run kubectl wait --for=condition=established --timeout=120s \
      crd/inferenceservices.serving.kserve.io

    log "helm install: kserve $KSERVE_VER (RawDeployment)"
    # NOTE the nesting on enableGatewayApi: gateway.ingressGateway.enableGatewayApi.
    # Helm silently ignores a wrong key rather than erroring.
    run helm upgrade --install kserve "oci://ghcr.io/kserve/charts/kserve" \
      --version "$KSERVE_VER" -n "$KSERVE_NS" \
      --set kserve.controller.deploymentMode=RawDeployment \
      --set kserve.controller.gateway.ingressGateway.enableGatewayApi=false
    # Controller readiness via kubectl, not helm --wait, for the same reason.
    run kubectl rollout status deploy -n "$KSERVE_NS" --timeout=600s
  else
    sync_app kserve-crd
    run kubectl wait --for=condition=established --timeout=120s \
      crd/inferenceservices.serving.kserve.io
    sync_app kserve
  fi
  info "KServe ready (RawDeployment mode)"
}

uncordon() {
  log "uncordon $NODE"
  # Must precede the engines. The device-plugin DaemonSet tolerates the
  # unschedulable taint automatically; the engine Deployments do not and would
  # sit Pending forever.
  run kubectl uncordon "$NODE"
}

engines() {
  if [[ $MODE == manual ]]; then
    log "kubectl apply -k k8s/apps/llm-serving"
    run kubectl apply -k "$REPO_ROOT/k8s/apps/llm-serving"
  else
    sync_app llm-serving
  fi

  log "waiting for InferenceServices (first start pulls ~60GB of weights)"
  run kubectl wait --for=condition=Ready --timeout=2400s \
    inferenceservice --all -n "$NS" \
    || info "not all Ready yet — check: kubectl get pods -n $NS"
}

smoke() {
  log "smoke test"
  $DRY && { info "[dry-run] would curl the three engines"; return 0; }

  # Engines are ClusterIP behind a default-deny NetworkPolicy, so probe from
  # inside the cluster with the llm-client label that the policy allows.
  #
  # KServe names the Service "<isvc>-predictor" (constants.PredictorServiceName)
  # on port 80 -> containerPort 8080 — NOT "<isvc>" on 8080.
  probe() {
    local isvc=$1 path=$2 body=$3
    printf '    %-24s ' "$isvc"
    if kubectl run "smoke-$RANDOM" -n "$NS" --rm -i --restart=Never -q \
         --image=curlimages/curl:8.11.1 \
         --labels="llm-client=true" \
         --command -- curl -sf -m 300 -X POST \
         "http://${isvc}-predictor.${NS}.svc.cluster.local/${path}" \
         -H 'Content-Type: application/json' -d "$body" >/dev/null 2>&1; then
      printf '\033[32mOK\033[0m\n'
    else
      printf '\033[31mFAIL\033[0m\n'
    fi
  }

  # Thinking models (reasoning, ornith) spend their first tokens inside a
  # <think> block, so a small max_tokens returns EMPTY content and looks like a
  # failure. Give them room.
  probe reasoning v1/completions \
    '{"model":"reasoning","prompt":"2+2=","max_tokens":200}'
  probe ornith    v1/chat/completions \
    '{"model":"ornith:35b","messages":[{"role":"user","content":"hi"}],"max_tokens":400}'
  # Same Ollama endpoint, second co-resident model on GPU2.
  probe ornith    v1/chat/completions \
    '{"model":"qwen2.5:3b-instruct","messages":[{"role":"user","content":"hi"}],"max_tokens":64}'
  # Reward model: scores via /pooling with an `input` field, NOT /v1/completions.
  probe verifier  pooling \
    '{"model":"verifier","input":"Step 1: 2+2=4."}'
  # TEI encoders (CPU): /embed and /rerank, not an OpenAI-shaped API.
  # Nomic v2 needs the search_document:/search_query: prefix from the client.
  probe embeddings embed \
    '{"inputs":"search_query: hello"}'
  probe reranker   rerank \
    '{"query":"how to sort a list","texts":["use sorted()","unrelated text"]}'
}

status() {
  log "status"
  kubectl get node "$NODE" -o wide
  echo
  kubectl get inferenceservice -n "$NS" 2>/dev/null || info "no InferenceServices yet"
  echo
  kubectl get pods -n "$NS" -o wide 2>/dev/null || true
  echo
  info "GPU allocatable: $(kubectl get node "$NODE" -o jsonpath='{.status.allocatable.nvidia\.com/gpu}' 2>/dev/null || echo none)"
}

# Every check that this bring-up has actually needed, in one command, so
# troubleshooting never requires ad-hoc kubectl archaeology again. Read-only.
doctor() {
  log "doctor — full diagnostic"
  set +e   # advisory only: never abort on an absent resource
  local fail=0
  chk() { # chk <label> <expected> <actual>
    if [[ "$2" == "$3" ]]; then printf '    \033[32m✓\033[0m %-38s %s\n' "$1" "$3"
    else printf '    \033[31m✗\033[0m %-38s got=%s want=%s\n' "$1" "${3:-<empty>}" "$2"; fail=$((fail+1)); fi
  }

  echo "  [node]"
  chk "Ready" True "$(kubectl get node "$NODE" -o jsonpath='{.status.conditions[?(@.type=="Ready")].status}' 2>/dev/null)"
  local cordon; cordon=$(kubectl get node "$NODE" -o jsonpath='{.spec.unschedulable}' 2>/dev/null || true)
  printf '    %-40s %s\n' "cordoned" "${cordon:-false} (must be false before engines)"
  chk "driver extension label" present \
    "$(kubectl get node "$NODE" -o jsonpath='{.metadata.labels}' 2>/dev/null | grep -q nonfree-kmod-nvidia && echo present)"
  chk "label nvidia.com/gpu" true "$(kubectl get node "$NODE" -o jsonpath='{.metadata.labels.nvidia\.com/gpu}' 2>/dev/null)"
  chk "allocatable nvidia.com/gpu" 4 "$(kubectl get node "$NODE" -o jsonpath='{.status.allocatable.nvidia\.com/gpu}' 2>/dev/null)"

  echo "  [containerd / runtime]"
  # The Talos extension registers the handler; the RuntimeClass object is ours.
  chk "RuntimeClass nvidia" nvidia "$(kubectl get runtimeclass nvidia -o jsonpath='{.handler}' 2>/dev/null)"

  echo "  [device plugin]"
  local st; st=$(helm list -n "$GPU_NS" -a -o json 2>/dev/null \
    | python3 -c 'import json,sys;print(next((r["status"] for r in json.load(sys.stdin) if r["name"]=="nvidia-device-plugin"),""))' 2>/dev/null || true)
  chk "helm release status" deployed "$st"
  [[ $st == pending-* || $st == failed ]] && \
    info "  -> stuck release: run '$0 unstick' (helm refuses upgrades in this state)"
  chk "DaemonSet desired" 1 "$(kubectl get ds -n "$GPU_NS" nvidia-device-plugin -o jsonpath='{.status.desiredNumberScheduled}' 2>/dev/null)"
  chk "DaemonSet ready"   1 "$(kubectl get ds -n "$GPU_NS" nvidia-device-plugin -o jsonpath='{.status.numberReady}' 2>/dev/null)"
  local aff; aff=$(kubectl get ds -n "$GPU_NS" nvidia-device-plugin -o jsonpath='{.spec.template.spec.affinity}' 2>/dev/null || true)
  [[ -z $aff ]] && printf '    \033[32m✓\033[0m %-38s removed\n' "chart nodeAffinity (needs NFD)" \
                || { printf '    \033[31m✗\033[0m %-38s present — requires NFD labels, will match 0 nodes\n' "chart nodeAffinity"; fail=$((fail+1)); }
  local ev; ev=$(kubectl get events -n "$GPU_NS" --field-selector reason=FailedCreate \
    -o jsonpath='{.items[-1:].message}' 2>/dev/null || true)
  [[ -n $ev ]] && info "  last FailedCreate: ${ev:0:110}"

  echo "  [kserve]"
  chk "InferenceService CRD" true \
    "$(kubectl get crd inferenceservices.serving.kserve.io >/dev/null 2>&1 && echo true)"
  local dm; dm=$(kubectl get cm inferenceservice-config -n "$KSERVE_NS" -o jsonpath='{.data.deploy}' 2>/dev/null \
    | grep -o '"defaultDeploymentMode": *"[^"]*"' | sed 's/.*"\([^"]*\)"$/\1/' || true)
  chk "deploymentMode" RawDeployment "$dm"

  echo "  [engines]"
  if kubectl get ns "$NS" >/dev/null 2>&1; then
    kubectl get inferenceservice -n "$NS" --no-headers 2>/dev/null | awk '{printf "    %-24s ready=%s\n",$1,$3}'
    kubectl get pods -n "$NS" --no-headers 2>/dev/null \
      | awk '{printf "    %-40s %s %s\n",$1,$3,$5}'
  else
    info "namespace $NS absent (engines not deployed yet)"
  fi

  echo
  [[ $fail -eq 0 ]] && info "all checks passed" || info "$fail check(s) failed"
  set -e
  return 0
}

# Clear a helm release wedged in pending-* / failed. Helm blocks every
# subsequent upgrade with "another operation is in progress" until this is done.
unstick() {
  log "unstick helm release"
  local st; st=$(helm list -n "$GPU_NS" -a -o json 2>/dev/null \
    | python3 -c 'import json,sys;print(next((r["status"] for r in json.load(sys.stdin) if r["name"]=="nvidia-device-plugin"),""))' 2>/dev/null)
  info "current status: ${st:-<no release>}"
  case $st in
    deployed) info "nothing to do" ;;
    "")       info "no release; nothing to do" ;;
    *)        # Roll back to the last deployed revision; if none, uninstall.
              if helm history nvidia-device-plugin -n "$GPU_NS" 2>/dev/null | grep -q deployed; then
                run helm rollback nvidia-device-plugin -n "$GPU_NS"
              else
                run helm uninstall nvidia-device-plugin -n "$GPU_NS"
              fi ;;
  esac
  # The DaemonSet controller backs off after repeated FailedCreate and will not
  # retry for many minutes even after the underlying cause is fixed. Force it.
  if kubectl get ds -n "$GPU_NS" nvidia-device-plugin >/dev/null 2>&1; then
    log "nudging DaemonSet (clears admission backoff)"
    run kubectl rollout restart ds/nvidia-device-plugin -n "$GPU_NS"
  fi
}

# Remove everything this script creates, so a retry starts from a clean slate.
# Does NOT touch the node's Talos config or the models PVC by default.
teardown() {
  log "teardown"
  info "this removes: llm-serving ns, kserve, device plugin, RuntimeClass"
  if ! $DRY; then
    printf '    type "yes" to proceed: '; local a; read -r a
    [[ $a == yes ]] || die "aborted"
  fi
  run kubectl delete -k "$REPO_ROOT/k8s/apps/llm-serving" --ignore-not-found
  run helm uninstall kserve -n "$KSERVE_NS" --ignore-not-found 2>/dev/null || true
  run helm uninstall kserve-crd -n "$KSERVE_NS" --ignore-not-found 2>/dev/null || true
  run helm uninstall nvidia-device-plugin -n "$GPU_NS" --ignore-not-found 2>/dev/null || true
  run kubectl delete -k "$REPO_ROOT/k8s/apps/gpu-runtimeclass" --ignore-not-found
  info "PVC llm-models in ns $NS was NOT deleted (holds ~56Gi of downloaded weights)"
  info "delete explicitly if you want a cold re-download:"
  info "  kubectl delete pvc llm-models -n $NS"
}

all() { preflight; gpu_plugin; kserve; uncordon; engines; smoke; status; }

# ---------------------------------------------------------------- main

STAGE=all
for a in "$@"; do
  case $a in
    --manual)  MODE=manual ;;
    --argocd)  MODE=argocd ;;
    --dry-run) DRY=true ;;
    preflight|runtimeclass|gpu-plugin|kserve|uncordon|engines|smoke|status|doctor|unstick|teardown|all) STAGE=$a ;;
    -h|--help) sed -n '2,28p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    *) die "unknown argument: $a" ;;
  esac
done

$DRY && info "DRY RUN — no changes will be made"
case $STAGE in
  preflight)  preflight ;;
  runtimeclass) runtimeclass ;;
  gpu-plugin) preflight; gpu_plugin ;;
  kserve)     preflight; kserve ;;
  uncordon)   uncordon ;;
  engines)    preflight; engines ;;
  smoke)      smoke ;;
  status)     status ;;
  doctor)     doctor ;;
  unstick)    unstick ;;
  teardown)   teardown ;;
  all)        all ;;
esac

log "done"