#!/usr/bin/env bash
# ============================================================================
# spark_api.sh — Spark-X2.5-4B OpenAI 兼容 API 服务管理（Colab T4 版）
# 位置: /content/drive/MyDrive/spark-t4/spark_api.sh
#
# 能力:
#   · OpenAI 兼容端点 /v1/chat/completions /v1/completions /v1/models
#   · 流式输出: 请求体加 "stream": true 即可（SSE，llama-server 原生支持）
#   · 工具调用: --jinja 启用模型内置 chat template 的 function calling
#   · 大上下文: CTX 环境变量控制，默认 131072 (128k)，最大可设 1048576 (1M)
#     —— 注意 T4 只有 16GB 显存，CTX 越大 KV cache 越大；1M 需配合 KV 量化
#     (KVQ=q8_0 或 q4_0) 才可能放下，且长 prompt 预填充极慢。启动日志会打印
#     KV/计算缓冲大小，若 OOM 请下调 CTX。
#   · 鉴权: 默认自动生成 API Key 存 Drive 并打印；API_KEY=off 可关闭
#
# 用法:
#   bash spark_api.sh start     # 起服务（默认 Q8_0+1.7B草稿投机解码 @ CTX=262144, 思考默认关）
#   关闭投机解码: SPEC_TYPE=none CTX=500000 bash spark_api.sh start
#   开启思考模式: THINKING=on bash spark_api.sh start（⚠️ 非流式请求可能超网关 120s 超时→504）
#   bash spark_api.sh stop
#   bash spark_api.sh status
#   bash spark_api.sh test      # 非流式 + 流式各测一次
#   自定义: CTX=262144 KVQ=q8_0 bash spark_api.sh start
# ============================================================================
set -euo pipefail

BASE_DIR="${BASE_DIR:-/content/drive/MyDrive/spark-t4}"
MODEL_DIR="${MODEL_DIR:-/content/spark-t4-models}"
QUANT="${QUANT:-Q8_0}"                     # 默认 Q8_0(质量最佳+草稿接受率0.89≈52tok/s); 4bit省显存: QUANT=Q4_K_M
CTX="${CTX:-262144}"                       # 256k 默认（配合默认投机解码）；关草稿后可设 524288/500000(需KVQ)
PORT="${PORT:-8100}"                       # 8100 默认(Colab 的 8080 被占用); 本机另跑时 PORT=8080 覆盖
NGL="${NGL:-999}"
KVQ="${KVQ:-}"                             # KV cache 量化: q8_0 / q4_0，压显存换长上下文
API_KEY="${API_KEY:-auto}"                 # auto=首次生成并存Drive; off=不鉴权
THREADS="${THREADS:-2}"
NP="${NP:--}"                              # 并行槽位; -=用fork默认; 单槽1M用 NP=1
BATCH="${BATCH:--}"                        # 逻辑batch; 冲1M时建议 512
UBATCH="${UBATCH:--}"                      # 物理batch(预填充粒度); 冲1M时建议 512
EXTRA_ARGS="${EXTRA_ARGS:-}"               # 额外 llama-server 参数, 如 "-fa on"
# ---- 投机解码（默认启用! 实测 T4: Q8_0主模型+1.7B草稿 37→52.3 tok/s +41%, 接受率 0.891）----
# 关闭方式: SPEC_TYPE=none bash spark_api.sh start （或 DRAFT_MODEL= 置空）
# ⚠️ 草稿 KV/计算缓冲继承主 CTX: CTX>300000 会 OOM（实测 500k 爆显存），开草稿请保持 CTX≤262144
# 草稿模型手动补下: hf download gasschina/Spark-X2.5-4B-build-cpp Spark-X2.5-1.7B-Q8_0.gguf --local-dir $MODEL_DIR
SPEC_TYPE="${SPEC_TYPE:-draft-simple}"     # 可选: draft-simple / ngram-simple / none
DRAFT_MODEL="${DRAFT_MODEL:-/content/spark-t4-models/Spark-X2.5-1.7B-Q8_0.gguf}"  # 默认 1.7B 草稿; 置空或 SPEC_TYPE=none 关闭
DRAFT_NMAX="${DRAFT_NMAX:-8}"              # 草稿步长上限(实测 8→16 无增益)
# ---- 思考模式默认关（2026-09-06 实测: aitun 隧道非流式请求 120s 硬超时→504）----
# 开思考后首响应延长 30~90s, aicq.me 等平台网关直接 504; 需要深度思考时 THINKING=on 重启
# 或单请求覆盖: 请求体加 "chat_template_kwargs":{"enable_thinking":true}
THINKING="${THINKING:-off}"                # off=服务端默认不思考(推荐); on=模型默认行为
MODEL_FILE="$MODEL_DIR/Spark-X2.5-4B-${QUANT}.gguf"
BIN="$BASE_DIR/llama.cpp-spark/build/bin/llama-server"
KEY_FILE="$BASE_DIR/.api_key"
LOG_FILE="$BASE_DIR/api-server.log"
PID_FILE="$BASE_DIR/api-server.pid"

GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; CYAN='\033[0;36m'; NC='\033[0m'
info() { echo -e "${GREEN}[API]${NC} $*"; }
warn() { echo -e "${YELLOW}[API]${NC} $*"; }
fail() { echo -e "${RED}[API]${NC} $*"; exit 1; }

resolve_api_key() {
    case "$API_KEY" in
        off)  API_KEY_VALUE="" ;;
        auto)
            if [[ -s "$KEY_FILE" ]]; then
                API_KEY_VALUE="$(cat "$KEY_FILE")"
            else
                API_KEY_VALUE="sk-$(head -c 24 /dev/urandom | base64 | tr -dc 'a-zA-Z0-9' | head -c 32)"
                echo "$API_KEY_VALUE" > "$KEY_FILE"
                chmod 600 "$KEY_FILE" 2>/dev/null || true
            fi ;;
        *)    API_KEY_VALUE="$API_KEY" ;;
    esac
}

do_start() {
    [[ -x "$BIN" ]] || fail "llama-server 不存在: $BIN （先跑 deploy_spark_t4.sh build）"
    [[ -f "$MODEL_FILE" ]] || fail "模型不存在: $MODEL_FILE （先跑 deploy_spark_t4.sh download）"
    do_stop 2>/dev/null || true
    resolve_api_key

    local kv_args=""
    [[ -n "$KVQ" ]] && kv_args="-ctk $KVQ -ctv $KVQ"
    local np_args=""
    [[ "$NP" != "-" ]] && np_args="-np $NP"
    local batch_args=""
    [[ "$BATCH" != "-" ]] && batch_args="-b $BATCH"
    [[ "$UBATCH" != "-" ]] && batch_args="$batch_args -ub $UBATCH"
    local key_args=""
    [[ -n "$API_KEY_VALUE" ]] && key_args="--api-key $API_KEY_VALUE"
    local spec_args=""
    if [[ "$SPEC_TYPE" != "none" && -n "$DRAFT_MODEL" ]]; then
        if [[ -f "$DRAFT_MODEL" ]]; then
            if (( CTX > 300000 )); then
                warn "CTX=$CTX > 300k 时草稿计算缓冲会 OOM（实测 500k 爆显存），建议 CTX≤262144 或 SPEC_TYPE=none"
            fi
            spec_args="--spec-type $SPEC_TYPE -md $DRAFT_MODEL -ngld 999 -ctkd ${DKVQ:-q4_0} -ctvd ${DKVQ:-q4_0} --spec-draft-n-max $DRAFT_NMAX --spec-draft-n-min 2"
            info "投机解码已启用: $SPEC_TYPE 草稿=$DRAFT_MODEL n_max=$DRAFT_NMAX"
            if [[ "$QUANT" == Q4* ]]; then
                warn "4bit 权重挂草稿实测反而慢一半(27 vs 48.75 tok/s, 接受率仅0.37-0.61)——草稿仅建议 Q8_0 使用"
            fi
        else
            warn "草稿模型不存在($DRAFT_MODEL)，本次无草稿启动。补下: hf download gasschina/Spark-X2.5-4B-build-cpp Spark-X2.5-1.7B-Q8_0.gguf --local-dir $MODEL_DIR"
        fi
    fi
    local think_args=""
    if [[ "$THINKING" == "off" ]]; then
        think_args='--chat-template-kwargs {"enable_thinking":false}'
        info "思考模式: 服务端默认关闭（避免非流式请求超网关 120s 限制→504）; 单请求可传 enable_thinking:true 恢复"
    fi

    info "启动: quant=$QUANT ctx=$CTX port=$PORT kv_quant=${KVQ:-无} auth=$([[ -n "$API_KEY_VALUE" ]] && echo ON || echo OFF)"
    nohup "$BIN" \
        -m "$MODEL_FILE" \
        -ngl "$NGL" \
        -c "$CTX" \
        -t "$THREADS" \
        --jinja \
        --host 0.0.0.0 \
        --port "$PORT" \
        $kv_args $np_args $batch_args $key_args $spec_args $think_args $EXTRA_ARGS \
        > "$LOG_FILE" 2>&1 &
    echo $! > "$PID_FILE"

    # 等待模型加载完毕（大模型 + 长上下文分配 KV 需要时间）
    info "等待模型加载..."
    for i in $(seq 1 60); do
        sleep 5
        if ! kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
            fail "进程退出，日志尾部:\n$(tail -25 "$LOG_FILE")"
        fi
        if curl -s --max-time 3 "http://127.0.0.1:$PORT/health" 2>/dev/null | grep -q '"status"'; then
            info "服务就绪!"
            echo -e "${CYAN}──────────────────────────────────────────────────${NC}"
            echo -e "  本机地址:  http://127.0.0.1:$PORT/v1/chat/completions"
            echo -e "  API Key:   ${API_KEY_VALUE:-（未启用鉴权）}"
            echo -e "  上下文:    $CTX tokens"
            [[ -n "$spec_args" ]] && echo -e "  投机解码:  $SPEC_TYPE + 1.7B草稿 n_max=$DRAFT_NMAX（实测 +41%）"
            echo -e "  思考模式:  $([[ "$THINKING" == "off" ]] && echo '默认关闭(504 防护)' || echo '默认开启')"
            echo -e "  流式输出:  请求体加 \"stream\": true"
            echo -e "  模型名:    任意值均可（服务端忽略 model 字段）"
            echo -e "${CYAN}──────────────────────────────────────────────────${NC}"
            # 打印显存占用摘要
            nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader | \
                awk '{print "  显存占用: " $0}'
            return 0
        fi
    done
    fail "300 秒未就绪，日志尾部:\n$(tail -25 "$LOG_FILE")"
}

do_stop() {
    if [[ -f "$PID_FILE" ]] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
        kill "$(cat "$PID_FILE")" && info "已停止 PID $(cat "$PID_FILE")"
    fi
    rm -f "$PID_FILE"
}

do_status() {
    if [[ -f "$PID_FILE" ]] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
        info "运行中 (PID $(cat "$PID_FILE"), ctx=$CTX)"
        curl -s --max-time 5 "http://127.0.0.1:$PORT/health" 2>/dev/null | head -c 300 || true; echo
        tail -3 "$LOG_FILE" 2>/dev/null || true
    else
        warn "未运行"
    fi
}

do_test() {
    resolve_api_key
    local auth=()
    [[ -n "$API_KEY_VALUE" ]] && auth=(-H "Authorization: Bearer $API_KEY_VALUE")
    info "== 非流式 =="
    curl -s --max-time 120 "http://127.0.0.1:$PORT/v1/chat/completions" \
        "${auth[@]}" -H "Content-Type: application/json" \
        -d '{"model":"spark","messages":[{"role":"user","content":"回复两个字：收到"}],"max_tokens":256,"temperature":1.0,"top_p":0.95}' \
        | python3 -c "
import json,sys
d=json.load(sys.stdin)
m=d['choices'][0]['message']
print('reasoning:', (m.get('reasoning_content') or '')[:80])
print('content  :', m.get('content',''))"
    info "== 流式 (SSE) =="
    curl -sN --max-time 120 "http://127.0.0.1:$PORT/v1/chat/completions" \
        "${auth[@]}" -H "Content-Type: application/json" \
        -d '{"model":"spark","messages":[{"role":"user","content":"从1数到8，用顿号分隔"}],"stream":true,"max_tokens":256,"temperature":1.0,"top_p":0.95}' \
        | python3 -c "
import sys, json
n = 0
for line in sys.stdin:
    line = line.strip()
    if line.startswith('data: ') and line != 'data: [DONE]':
        try:
            j = json.loads(line[6:])
            delta = j['choices'][0].get('delta', {})
            c = delta.get('content') or ''
            if c:
                n += 1
                sys.stdout.write(c); sys.stdout.flush()
        except Exception:
            pass
print()
print(f'--- 收到 {n} 个流式 chunk，SSE 正常 ---' if n else '!!! 未收到任何流式内容 !!!')"
}

case "${1:-help}" in
    start)  do_start ;;
    stop)   do_stop ;;
    status) do_status ;;
    test)   do_test ;;
    *)      grep '^#   ' "$0" | sed 's/^#   //' ;;
esac