#!/usr/bin/env bash
# vllmctl — control the vLLM serving stack through the router's admin API.
#
#   ./vllmctl status                  per-service state: awake / sleeping /
#                                     offloaded, activity, wake in progress
#   ./vllmctl up [MODEL]              force-wake MODEL now (no MODEL = show
#                                     status: all services are always running,
#                                     the router wakes them on demand)
#   ./vllmctl down [MODEL]            offload MODEL (sleep level 2, frees host
#                                     RAM); no MODEL = all services
#   ./vllmctl sleep [MODEL]           light sleep (level 1, weights stay in RAM)
#   ./vllmctl list                    models on disk + which ones this stack serves
#   ./vllmctl logs [-f] [N] [SVC]     docker compose logs (SVC: text|ocr|embed|router)
#   ./vllmctl restart [SVC]           docker compose restart (rare, manual)
#   ./vllmctl pull REPO [NAME]        download a HF model into MODEL_ROOT
#
# MODEL accepts the service key (text/ocr/embed), the model name
# (Qwen3.6-35B-A3B-FP8, OvisOCR2, Qwen3-Embedding-8B) or an alias,
# case-insensitively; the router resolves them all.
#
# Wake/sleep/idle is owned by the router (127.0.0.1:8010 admin listener);
# `idle-watch` no longer exists.  Public API: http://<host>:8000/v1.
#
# docker: this user is in the docker group, so plain `docker` works in fresh
# login shells.  In a session started before the group change, invoke this
# script as:  sg docker -c "./vllmctl pull <repo>"

set -u

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

C_GREEN=$'\e[32m'; C_RED=$'\e[31m'; C_YELLOW=$'\e[33m'; C_CYAN=$'\e[36m'; C_DIM=$'\e[2m'; C_OFF=$'\e[0m'
say()  { printf '%s\n' "$*"; }
ok()   { printf '%s%s%s\n' "$C_GREEN" "$*" "$C_OFF"; }
warn() { printf '%s%s%s\n' "$C_YELLOW" "$*" "$C_OFF"; }
err()  { printf '%s%s%s\n' "$C_RED" "$*" "$C_OFF" >&2; }

ADMIN="http://127.0.0.1:${VLLMCTL_ADMIN_PORT:-8010}"
PUBLIC_PORT="${VLLMCTL_PUBLIC_PORT:-8000}"
KEYS="text ocr embed"

# ---------------------------------------------------------------- config ----
env_get() {  # env_get KEY [default]
    local v
    v="$(grep -E "^$1=" "$ROOT/.env" 2>/dev/null | tail -n1 | cut -d= -f2-)"
    if [ -z "$v" ]; then printf '%s' "${2:-}"; else printf '%s' "$v"; fi
}

MODEL_ROOT="$(env_get MODEL_ROOT /data/home/renbaibing/huggingface)"

dcompose() {
    docker compose --progress plain --ansi never \
        --project-directory "$ROOT" -f "$ROOT/compose.yml" --env-file "$ROOT/.env" "$@"
}

# name -> docker compose SERVICE (for logs / restart only).
# NB: the router's compose service is `router` (container_name vllm-router);
# the vllm-* services are named after themselves.
docker_service_for() {
    case "$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')" in
        embed|embedding*|qwen3-embedding*) echo vllm-embed ;;
        ocr|ovis*)                         echo vllm-ocr ;;
        text|qwen*|chat|default)           echo vllm-text ;;
        router|vllm-router)                echo router ;;
        vllm-text|vllm-ocr|vllm-embed)     echo "$1" ;;
        *) return 1 ;;
    esac
}

# ---------------------------------------------------------------- admin -----
admin_alive() { curl -fs -m 3 "$ADMIN/health" >/dev/null 2>&1; }

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

admin_call() {  # admin_call METHOD PATH TIMEOUT -> sets HTTP_CODE and RESPONSE
    HTTP_CODE="$(curl -s -m "$3" -o "$RESPONSE_FILE" -w '%{http_code}' \
                     -X "$1" "$ADMIN$2" 2>/dev/null)"
    [ -n "$HTTP_CODE" ] || HTTP_CODE=000
    RESPONSE="$(cat "$RESPONSE_FILE")"
}

json_field() {  # json_field JSON PYEXPR -> value (python3 json)
    printf '%s' "$1" | python3 -c "
import json, sys
try:
    d = json.load(sys.stdin)
except Exception:
    sys.exit(1)
print(eval(sys.argv[1]))" "$2" 2>/dev/null
}

die_admin_down() {
    err "Router admin API not reachable at $ADMIN"
    err "Is the stack up?  Try: docker compose -f $ROOT/compose.yml ps"
    return 1
}

# ---------------------------------------------------------------- status ----
cmd_status() {
    admin_call GET /admin/status 10
    [ "$HTTP_CODE" = "200" ] || { die_admin_down; return 1; }

    # NB: the heredoc *is* stdin for `python3 -`, so the payload is passed as
    # a file argument rather than piped.
    python3 - "$PUBLIC_PORT" "$RESPONSE_FILE" <<'PY'
import json, sys

port, payload_path = sys.argv[1], sys.argv[2]
with open(payload_path) as fh:
    d = json.load(fh)
r = d["router"]
idle = r["idle"]
up = r.get("uptime_s")
up_txt = f"up {up:.0f}s" if up is not None and up < 3600 else f"up {up/3600:.1f}h"
tiers = f"idle tiers: sleep {idle['sleep_min']:g} min, offload {idle['offload_min']:g} min"
if not idle["enabled"]:
    tiers += "  (disabled)"

G, Y, D, R, OFF = "\033[32m", "\033[33m", "\033[2m", "\033[31m", "\033[0m"
print(f"router {G}{up_txt}{OFF}   {tiers}")
print(f"Public API : http://<this-host>:{port}/v1   (OpenAI-compatible, all 3 models)")
print(f"Admin API  : 127.0.0.1:{r['admin_port']}   (localhost only)")
print()
print(f"  {'SERVICE':<7} {'MODEL':<22} {'STATE':<11} {'ACTIVE':>6} {'IDLE':>8}  NOTES")
for key in ("text", "ocr", "embed"):
    s = d["services"][key]
    state = s["depth"] or "unknown"
    colour = {"awake": G, "offloaded": D}.get(state, Y)
    idle_s = s["last_activity_ago_s"]
    if idle_s is None:
        idle_txt = "-"
    elif idle_s >= 3600:
        idle_txt = f"{idle_s/3600:.1f}h"
    elif idle_s >= 60:
        idle_txt = f"{idle_s/60:.0f}m"
    else:
        idle_txt = f"{idle_s:.0f}s"
    notes = []
    if not s["reachable"]:
        notes.append(R + "unreachable" + OFF)
    if s.get("wake_recovery_pending"):
        notes.append(Y + "reload recovery pending" + OFF)
    if s["wake_in_progress"]:
        notes.append("waking...")
    if s["last_wake_latency_s"] is not None:
        notes.append(f"last wake {s['last_wake_latency_s']}s")
    if s["last_error"]:
        notes.append(s["last_error"])
    print(f"  {key:<7} {s['model']:<22} {colour}{state:<11}{OFF} "
          f"{s['active_requests']:>6} {idle_txt:>8}  {' '.join(notes)}")
PY
}

# -------------------------------------------------------------------- up ----
cmd_up() {
    local target="${1:-}" body key
    if [ -z "$target" ]; then
        say "All services are always running -- the router wakes them on demand."
        say "Nothing to do.  Use '${C_CYAN}./vllmctl up MODEL${C_OFF}' to force-wake one now."
        cmd_status
        return 0
    fi
    admin_alive || { die_admin_down; return 1; }
    # Level-2 wakes can take a minute or two (weights come back from NFS).
    admin_call POST "/admin/wake/$(uri_escape "$target")" 400
    local body="$RESPONSE"
    case "$HTTP_CODE" in
        200)
            local lat
            lat="$(json_field "$body" "d.get('latency_s')")"
            [ -n "$lat" ] && lat=" in ${lat}s"
            ok "Awake: $(json_field "$body" "d['model']") ($target)${lat:-}"
            return 0
            ;;
        404) err "Unknown model or service: '$target'"; return 1 ;;
        503)
            warn "'$target' is still waking ($(json_field "$body" "d['error']['sleep_depth']"))."
            say "  Retry-After: $(json_field "$body" "d['retry_after_s']")s   est. $(json_field "$body" "d['estimated_wake_seconds']")s"
            [ -n "$(json_field "$body" "d['error'].get('message','')")" ] && \
                say "  $(json_field "$body" "d['error']['message']")"
            return 1
            ;;
        *) err "Router returned HTTP $HTTP_CODE"; printf '%s\n' "$body"; return 1 ;;
    esac
}

uri_escape() { python3 -c "import urllib.parse,sys;print(urllib.parse.quote(sys.argv[1],safe=''))" "$1"; }

# ------------------------------------------------------------------ down ----
cmd_down() {
    local level="$1"; shift
    local targets="${1:-}" body rc=0 key
    admin_alive || { die_admin_down; return 1; }
    if [ -z "$targets" ]; then targets="$KEYS"; say "Offloading all services (sleep level 2)..."; fi
    for key in $targets; do
        admin_call POST "/admin/sleep/$(uri_escape "$key")?level=$level" 120
        local body="$RESPONSE"
        case "$HTTP_CODE" in
            200) ok "$(json_field "$body" "d['model']") ($key): $(sleep_state "$level")" ;;
            409)
                warn "$key: refused -- $(json_field "$body" "d.get('reason','busy')")"
                rc=1 ;;
            404) err "Unknown model or service: '$key'"; rc=1 ;;
            *)   err "$key: HTTP $HTTP_CODE $(printf '%s' "$body" | head -c 200)"; rc=1 ;;
        esac
    done
    return $rc
}

sleep_state() { [ "$1" = "2" ] && printf '%s' 'offloaded (weights freed)' || printf '%s' 'sleeping (weights in RAM)'; }

# ------------------------------------------------------------------ list ----
cmd_list() {
    say "${C_CYAN}Available models in $MODEL_ROOT${C_OFF}"
    local served d name found=0
    served="$(curl -fs -m 3 "http://127.0.0.1:$PUBLIC_PORT/v1/models" 2>/dev/null)"
    for d in "$MODEL_ROOT"/*/; do
        [ -d "$d" ] || continue
        name="$(basename "$d")"
        if compgen -G "$d/*.safetensors" >/dev/null || [ -f "$d/config.json" ]; then
            found=1
            local mark="${C_DIM}on disk${C_OFF}"
            [ -n "$served" ] && printf '%s' "$served" | grep -q "\"$name\"" && mark="${C_GREEN}*served${C_OFF}"
            local size
            size="$(du -sh "$d" 2>/dev/null | cut -f1)"
            printf '  %-40s %6s   %s\n' "$name" "$size" "$mark"
        fi
    done
    [ "$found" = 1 ] || warn '  (none found — pull one with: ./vllmctl pull <hf-repo>)'
}

# ------------------------------------------------------------------ logs ----
cmd_logs() {
    local follow='' n=100 svc=''
    for a in "$@"; do
        case "$a" in
            -f|--follow) follow='--follow' ;;
            *[!0-9]*) svc="$a" ;;
            *) n="$a" ;;
        esac
    done
    local args=($follow --tail "$n")
    if [ -n "$svc" ]; then
        local service
        service="$(docker_service_for "$svc")" || { err "Unknown service '$svc' (text|ocr|embed|router)"; return 1; }
        args+=("$service")
    fi
    dcompose logs "${args[@]}"
}

# --------------------------------------------------------------- restart ----
cmd_restart() {
    local svc="${1:-}" service
    [ -n "$svc" ] || { err 'Usage: ./vllmctl restart SERVICE (text|ocr|embed|router)'; return 1; }
    service="$(docker_service_for "$svc")" || { err "Unknown service '$svc'"; return 1; }
    warn "Restarting $service -- cold start can take 2-10 min (NFS weights)."
    dcompose restart "$service"
}

# ------------------------------------------------------------------ pull ----
cmd_pull() {
    local repo="${1:-}" name="${2:-}"
    [ -n "$repo" ] || { err 'Usage: ./vllmctl pull <hf-repo-id> [local-name]'; return 1; }
    [ -n "$name" ] || name="$(basename "$repo")"
    if [ -e "$MODEL_ROOT/$name" ]; then
        err "$MODEL_ROOT/$name already exists"
        return 1
    fi
    say "Downloading '$repo' → $MODEL_ROOT/$name  (Ctrl-C safe to abort)"
    local image="vllm/vllm-openai:$(env_get VLLM_VERSION v0.27.1)"
    docker run --rm -i \
        --name vllm-pull \
        -v "$MODEL_ROOT:/models" \
        --env-file "$ROOT/.env" \
        --entrypoint python3 \
        "$image" \
        -c "from huggingface_hub import snapshot_download; p=snapshot_download('$repo', local_dir='/models/$name'); print('done:', p)"
    local rc=$?
    [ $rc -eq 0 ] && ok "Pulled '$name'." || err "Pull failed (rc=$rc)"
    return $rc
}

# ------------------------------------------------------------------ main ----
case "${1:-help}" in
    status)   shift; cmd_status "$@" ;;
    up)       shift; cmd_up "$@" ;;
    down|offload)  shift; cmd_down 2 "$@" ;;
    sleep|nap)     shift; cmd_down 1 "$@" ;;
    stop)     shift; cmd_down 2 "$@" ;;
    list)     shift; cmd_list "$@" ;;
    logs)     shift; cmd_logs "$@" ;;
    restart)  shift; cmd_restart "$@" ;;
    pull)     shift; cmd_pull "$@" ;;
    help|--help|-h|*)
        sed -n '2,26p' "$0" | sed 's/^# \{0,1\}//'
        ;;
esac