#!/usr/bin/env bash
set -uo pipefail

CONTAINER="${VLLM_CONTAINER:-vllm}"
LOG_FILE="${VLLM_LOG_FILE:-/var/log/vllm.log}"
SOURCE="docker"
TAIL=""
SINCE="10m"

usage() {
    cat <<EOF
Usage: vllm-log [-f] [-c CONTAINER] [-n LINES] [-s SINCE]
  -f            read $LOG_FILE (tail -f) instead of docker logs
  -c CONTAINER  docker container name (default: $CONTAINER)
  -n LINES      start from last LINES (docker: --tail, file: -n). Disables -s.
  -s SINCE      docker logs --since (default: ${SINCE}. "" = full log)
  -h            help
EOF
}

while getopts "fc:n:s:h" opt; do
    case "$opt" in
        f) SOURCE="file" ;;
        c) CONTAINER="$OPTARG" ;;
        n) TAIL="$OPTARG"; SINCE="" ;;
        s) SINCE="$OPTARG"; TAIL="" ;;
        h) usage; exit 0 ;;
        *) usage >&2; exit 2 ;;
    esac
done

if [ "$SOURCE" = "file" ]; then
    [ -r "$LOG_FILE" ] || { echo "cannot read $LOG_FILE" >&2; exit 1; }
    if [ -n "$TAIL" ]; then
        tail -n "$TAIL" -f "$LOG_FILE"
    else
        tail -n 200 -f "$LOG_FILE"
    fi
else
    if ! docker inspect --format '{{.Name}}' "$CONTAINER" >/dev/null 2>&1; then
        echo "container not found: $CONTAINER" >&2
        exit 1
    fi
    ARGS=(logs -f)
    [ -n "$TAIL" ] && ARGS+=(--tail "$TAIL")
    [ -n "$SINCE" ] && ARGS+=(--since "$SINCE")
    docker "${ARGS[@]}" "$CONTAINER" 2>&1
fi | gawk '
function fmt_pct(v, got) {
    return got ? sprintf("%5.1f%%", v) : "  n/a"
}

function flush_pending() {
    if (!have) return
    printf "\033[37m[%s]\033[0m \033[36mPP\033[0m: \033[32m%6.1f t/s\033[0m | \033[31mTG\033[0m: \033[32m%6.1f t/s\033[0m | \033[33mK/V Cache\033[0m: \033[36m%6s\033[0m | \033[35mPfx Cache\033[0m: \033[36m%5.1f%%\033[0m | \033[34mDraft Acc\033[0m: \033[36m%6s\033[0m\n", \
        pending_ts, pending_pp, pending_tg, fmt_pct(pending_kv, got_kv), pending_pfx, fmt_pct(pending_draft, got_draft)
    fflush()
    have = 0
}

/Avg prompt throughput:/ {
    flush_pending()
    if (match($0, /INFO ([0-9-]{5} [0-9:]{8})/, t)) pending_ts = t[1]
    match($0, /Avg prompt throughput: ([0-9.]+)/, p);  pending_pp  = p[1] + 0
    match($0, /Avg generation throughput: ([0-9.]+)/, g); pending_tg = g[1] + 0
    if (match($0, /Prefix cache hit rate: ([0-9.]+)/, c)) pending_pfx = c[1] + 0
    if (match($0, /KV cache usage: ([0-9.]+)/, k)) {
        pending_kv = k[1] + 0
        got_kv = 1
    } else {
        got_kv = 0
    }
    if (match($0, /Avg Draft acceptance rate: ([0-9.]+)/, d)) {
        pending_draft = d[1] + 0
        got_draft = 1
    } else {
        got_draft = 0
    }
    have = 1
    next
}

/Avg Draft acceptance rate:/ {
    if (!have) next
    if (match($0, /INFO ([0-9-]{5} [0-9:]{8})/, t) && t[1] != pending_ts) next
    match($0, /Avg Draft acceptance rate: ([0-9.]+)/, d)
    pending_draft = d[1] + 0
    got_draft = 1
    flush_pending()
}

END { flush_pending() }
'