#!/bin/bash
set -euo pipefail

# ============================================================
#  vLLM Serving Script for Qwen3.6-35B-A3B-FP8
# ============================================================
#  All settings are overridable via environment variables.
#  See docker-compose.yaml for the defaults.
# ============================================================

# --- Model & Engine ---
MODEL_PATH="${MODEL_PATH:-/models/Qwen/Qwen3.6-35B-A3B-FP8}"
TP_SIZE="${TP_SIZE:-2}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.90}"
PORT="${PORT:-8000}"
SERVED_NAME="${SERVED_NAME:-qwen3.6-35b-a3b}"

# --- KV Cache Quantization ---
# int4_per_token_head: 4-bit per-token-head KV cache (4x smaller than bf16).
# Other options: auto, fp8, fp8_e5m2, fp8_e4m3, int8_per_token_head, nvfp4
KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-int4_per_token_head}"

# --- Speculative Decoding / Multi-Token Prediction (MTP) ---
# "mtp" enables the model's built-in MTP head for speculative decoding.
# The model has mtp_num_hidden_layers=1, so 1 speculative token is used.
# Set SPEC_METHOD to empty string to disable MTP.
SPEC_METHOD="${SPEC_METHOD:-mtp}"
SPEC_TOKENS="${SPEC_TOKENS:-1}"

# --- CUDA Graph Capture ---
MAX_CUDAGRAPH_CAPTURE_SIZE="${MAX_CUDAGRAPH_CAPTURE_SIZE:-32}"

# --- Max Concurrent Sequences ---
MAX_NUM_SEQS="${MAX_NUM_SEQS:-32}"

# --- Max Batched Tokens ---
MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-4096}"

# --- Enforce Eager (disable CUDA graphs) ---
# 0 = enable CUDA graphs (recommended for performance).
# 1 = disable CUDA graphs (fallback if OOM persists).
ENFORCE_EAGER="${ENFORCE_EAGER:-0}"

# --- API Authentication & Security ---
# Comma-separated list of API keys. If set, clients must send one in
# the Authorization header: "Bearer <key>". Leave empty to disable auth.
API_KEY="${API_KEY:-}"

# --- CORS (Cross-Origin Resource Sharing) ---
# JSON array strings, e.g. ALLOWED_ORIGINS='["https://app.example.com"]'
ALLOWED_ORIGINS="${ALLOWED_ORIGINS:-}"
ALLOWED_METHODS="${ALLOWED_METHODS:-}"
ALLOWED_HEADERS="${ALLOWED_HEADERS:-}"
# 1 = allow credentials, 0 = disallow (default).
ALLOW_CREDENTIALS="${ALLOW_CREDENTIALS:-0}"

# --- SSL / TLS ---
SSL_KEYFILE="${SSL_KEYFILE:-}"
SSL_CERTFILE="${SSL_CERTFILE:-}"
SSL_CA_CERTS="${SSL_CA_CERTS:-}"
SSL_CERT_REQS="${SSL_CERT_REQS:-}"
SSL_CIPHERS="${SSL_CIPHERS:-}"
ENABLE_SSL_REFRESH="${ENABLE_SSL_REFRESH:-0}"

# --- Chat Template & Tool Calling ---
CHAT_TEMPLATE="${CHAT_TEMPLATE:-}"
CHAT_TEMPLATE_CONTENT_FORMAT="${CHAT_TEMPLATE_CONTENT_FORMAT:-}"
TRUST_REQUEST_CHAT_TEMPLATE="${TRUST_REQUEST_CHAT_TEMPLATE:-0}"
DEFAULT_CHAT_TEMPLATE_KWARGS="${DEFAULT_CHAT_TEMPLATE_KWARGS:-}"
RESPONSE_ROLE="${RESPONSE_ROLE:-}"
RETURN_TOKENS_AS_TOKEN_IDS="${RETURN_TOKENS_AS_TOKEN_IDS:-0}"
ENABLE_AUTO_TOOL_CHOICE="${ENABLE_AUTO_TOOL_CHOICE:-0}"
TOOL_CALL_PARSER="${TOOL_CALL_PARSER:-}"
TOOL_PARSER_PLUGIN="${TOOL_PARSER_PLUGIN:-}"
TOOL_SERVER="${TOOL_SERVER:-}"
EXCLUDE_TOOLS_WHEN_TOOL_CHOICE_NONE="${EXCLUDE_TOOLS_WHEN_TOOL_CHOICE_NONE:-0}"

# --- Logging & Diagnostics ---
UVICORN_LOG_LEVEL="${UVICORN_LOG_LEVEL:-}"
DISABLE_UVICORN_ACCESS_LOG="${DISABLE_UVICORN_ACCESS_LOG:-0}"
DISABLE_ACCESS_LOG_FOR_ENDPOINTS="${DISABLE_ACCESS_LOG_FOR_ENDPOINTS:-}"
MAX_LOG_LEN="${MAX_LOG_LEN:-}"
ENABLE_LOG_REQUESTS="${ENABLE_LOG_REQUESTS:-0}"
ENABLE_LOG_OUTPUTS="${ENABLE_LOG_OUTPUTS:-0}"
ENABLE_LOG_DELTAS="${ENABLE_LOG_DELTAS:-}"
LOG_ERROR_STACK="${LOG_ERROR_STACK:-0}"
LOG_CONFIG_FILE="${LOG_CONFIG_FILE:-}"

# --- API Features ---
ROOT_PATH="${ROOT_PATH:-}"
ENABLE_PROMPT_TOKENS_DETAILS="${ENABLE_PROMPT_TOKENS_DETAILS:-0}"
ENABLE_PER_REQUEST_METRICS="${ENABLE_PER_REQUEST_METRICS:-0}"
ENABLE_SERVER_LOAD_TRACKING="${ENABLE_SERVER_LOAD_TRACKING:-0}"
ENABLE_FORCE_INCLUDE_USAGE="${ENABLE_FORCE_INCLUDE_USAGE:-0}"
ENABLE_TOKENIZER_INFO_ENDPOINT="${ENABLE_TOKENIZER_INFO_ENDPOINT:-0}"
DISABLE_FASTAPI_DOCS="${DISABLE_FASTAPI_DOCS:-0}"
ENABLE_REQUEST_ID_HEADERS="${ENABLE_REQUEST_ID_HEADERS:-0}"
ENABLE_OFFLINE_DOCS="${ENABLE_OFFLINE_DOCS:-0}"
H11_MAX_INCOMPLETE_EVENT_SIZE="${H11_MAX_INCOMPLETE_EVENT_SIZE:-}"
H11_MAX_HEADER_COUNT="${H11_MAX_HEADER_COUNT:-}"
FINGERPRINT_MODE="${FINGERPRINT_MODE:-}"
FINGERPRINT_VALUE="${FINGERPRINT_VALUE:-}"

echo "============================================"
echo "  Qwen3.6-35B-A3B-FP8 vLLM Server"
echo "============================================"
echo "  Model:              ${MODEL_PATH}"
echo "  Tensor Parallel:     ${TP_SIZE}"
echo "  Max Model Length:    ${MAX_MODEL_LEN}"
echo "  GPU Memory Util:    ${GPU_MEM_UTIL}"
echo "  Port:               ${PORT}"
echo "  Served Model Name:  ${SERVED_NAME}"
echo "  KV Cache Dtype:    ${KV_CACHE_DTYPE}"
echo "  Spec Method (MTP):  ${SPEC_METHOD}"
echo "  Spec Tokens:        ${SPEC_TOKENS}"
echo "  Max CUDA Graph:    ${MAX_CUDAGRAPH_CAPTURE_SIZE}"
echo "  Max Batched Tokens: ${MAX_NUM_BATCHED_TOKENS}"
echo "  Max Num Seqs:       ${MAX_NUM_SEQS}"
echo "  Enforce Eager:      ${ENFORCE_EAGER}"
echo "  API Key:           ${API_KEY:+(set)}"
echo "  CORS Origins:      ${ALLOWED_ORIGINS:-(*)}"
echo "  SSL Cert:          ${SSL_CERTFILE:-(none)}"
echo "  Chat Template:      ${CHAT_TEMPLATE:-(default)}"
echo "  Tool Choice:        ${ENABLE_AUTO_TOOL_CHOICE}"
echo "  Uvicorn Log:        ${UVICORN_LOG_LEVEL:-info}"
echo "============================================"
echo ""

# GPU visibility (set by CUDA_VISIBLE_DEVICES in docker-compose)
echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set}"
echo ""

# --- Helper: append flag only if env var is non-empty ---
add_if_set() {
    local var="$1"
    local flag="$2"
    local val="${!var}"
    if [ -n "${val}" ]; then
        CMD_ARGS+=("${flag}" "${val}")
    fi
}

# --- Helper: append boolean flag if env var is 1/true ---
add_bool() {
    local var="$1"
    local flag="$2"
    local val="${!var}"
    if [ "${val}" = "1" ] || [ "${val}" = "true" ]; then
        CMD_ARGS+=("${flag}")
    fi
}

# --- Build vllm serve command flags ---
CMD_ARGS=(
    --host 0.0.0.0
    --port "${PORT}"
    --tensor-parallel-size "${TP_SIZE}"
    --max-model-len "${MAX_MODEL_LEN}"
    --gpu-memory-utilization "${GPU_MEM_UTIL}"
    --served-model-name "${SERVED_NAME}"
    --trust-remote-code
    --enable-prefix-caching
    --kv-cache-dtype "${KV_CACHE_DTYPE}"
    --max-num-batched-tokens "${MAX_NUM_BATCHED_TOKENS}"
    --max-num-seqs "${MAX_NUM_SEQS}"
)

# Speculative decoding / MTP
if [ -n "${SPEC_METHOD}" ]; then
    CMD_ARGS+=(--spec-method "${SPEC_METHOD}")
    if [ -n "${SPEC_TOKENS}" ]; then
        CMD_ARGS+=(--spec-tokens "${SPEC_TOKENS}")
    fi
fi

# CUDA graph configuration
if [ "${ENFORCE_EAGER}" = "1" ]; then
    CMD_ARGS+=(--enforce-eager)
else
    CMD_ARGS+=(--max-cudagraph-capture-size "${MAX_CUDAGRAPH_CAPTURE_SIZE}")
fi

# --- API Authentication ---
add_if_set  API_KEY              --api-key

# --- CORS ---
add_if_set  ALLOWED_ORIGINS      --allowed-origins
add_if_set  ALLOWED_METHODS       --allowed-methods
add_if_set  ALLOWED_HEADERS       --allowed-headers
add_bool    ALLOW_CREDENTIALS    --allow-credentials

# --- SSL / TLS ---
add_if_set  SSL_KEYFILE           --ssl-keyfile
add_if_set  SSL_CERTFILE           --ssl-certfile
add_if_set  SSL_CA_CERTS          --ssl-ca-certs
add_if_set  SSL_CERT_REQS         --ssl-cert-reqs
add_if_set  SSL_CIPHERS           --ssl-ciphers
add_bool    ENABLE_SSL_REFRESH    --enable-ssl-refresh

# --- Chat Template & Tool Calling ---
add_if_set  CHAT_TEMPLATE             --chat-template
add_if_set  CHAT_TEMPLATE_CONTENT_FORMAT  --chat-template-content-format
add_bool    TRUST_REQUEST_CHAT_TEMPLATE  --trust-request-chat-template
add_if_set  DEFAULT_CHAT_TEMPLATE_KWARGS    --default-chat-template-kwargs
add_if_set  RESPONSE_ROLE              --response-role
add_bool    RETURN_TOKENS_AS_TOKEN_IDS  --return-tokens-as-token-ids
add_bool    ENABLE_AUTO_TOOL_CHOICE    --enable-auto-tool-choice
add_if_set  TOOL_CALL_PARSER          --tool-call-parser
add_if_set  TOOL_PARSER_PLUGIN        --tool-parser-plugin
add_if_set  TOOL_SERVER                --tool-server
add_bool    EXCLUDE_TOOLS_WHEN_TOOL_CHOICE_NONE  --exclude-tools-when-tool-choice-none

# --- Logging & Diagnostics ---
add_if_set  UVICORN_LOG_LEVEL                --uvicorn-log-level
add_bool    DISABLE_UVICORN_ACCESS_LOG       --disable-uvicorn-access-log
add_if_set  DISABLE_ACCESS_LOG_FOR_ENDPOINTS  --disable-access-log-for-endpoints
add_if_set  MAX_LOG_LEN                       --max-log-len
add_bool    ENABLE_LOG_REQUESTS               --enable-log-requests
add_bool    ENABLE_LOG_OUTPUTS                 --enable-log-outputs
if [ -n "${ENABLE_LOG_DELTAS:-}" ]; then
    if [ "${ENABLE_LOG_DELTAS}" = "0" ] || [ "${ENABLE_LOG_DELTAS}" = "false" ]; then
        CMD_ARGS+=(--disable-log-deltas)
    fi
fi
add_bool    LOG_ERROR_STACK                   --log-error-stack
add_if_set  LOG_CONFIG_FILE                    --log-config-file

# --- API Features ---
add_if_set  ROOT_PATH                          --root-path
add_bool    ENABLE_PROMPT_TOKENS_DETAILS        --enable-prompt-tokens-details
add_bool    ENABLE_PER_REQUEST_METRICS         --enable-per-request-metrics
add_bool    ENABLE_SERVER_LOAD_TRACKING         --enable-server-load-tracking
add_bool    ENABLE_FORCE_INCLUDE_USAGE          --enable-force-include-usage
add_bool    ENABLE_TOKENIZER_INFO_ENDPOINT       --enable-tokenizer-info-endpoint
add_bool    DISABLE_FASTAPI_DOCS                --disable-fastapi-docs
add_bool    ENABLE_REQUEST_ID_HEADERS           --enable-request-id-headers
add_bool    ENABLE_OFFLINE_DOCS                 --enable-offline-docs
add_if_set  H11_MAX_INCOMPLETE_EVENT_SIZE       --h11-max-incomplete-event-size
add_if_set  H11_MAX_HEADER_COUNT                --h11-max-header-count
add_if_set  FINGERPRINT_MODE                   --fingerprint-mode
add_if_set  FINGERPRINT_VALUE                   --fingerprint-value

exec vllm serve "${MODEL_PATH}" "${CMD_ARGS[@]}"