#!/usr/bin/env bash
# ============================================================================
# oneclick_start.sh — Spark-X2.5-4B Colab T4 一键启动（实战验证版 v2）
# 位置: /content/drive/MyDrive/spark-t4/oneclick_start.sh (Drive 持久化)
#
# 用法（Colab 任意 cell）:
#   !bash /content/drive/MyDrive/spark-t4/oneclick_start.sh
#
# 默认配置 v6（2026-09-06, 改回 Q8_0+草稿; 4bit+草稿实测反优化 27 tok/s 后回退）:
#   Q8_0 权重 + 1.7B 草稿投机解码 + CTX=262144(256k) + KV q4_0 + fa on + 别名 spark-x2.5-4b + 思考默认关
#   实测: 显存 11.8GB/15GB, 生成 51-52 tok/s, 草稿接受率 0.89, 质量最佳
#   4bit 档(QUANT=Q4_K_M): 自动关草稿裸跑 48.75 tok/s, 显存 6.6GB — 4bit 噪声→接受率 0.37-0.61, 挂草稿反而 27 tok/s
#
# 换配置（示例）:
#   高质量+草稿: QUANT=Q8_0 !bash .../oneclick_start.sh (草稿@256k, 52.3 tok/s, 11.8GB)
#   4bit 长上下文: CTX=500000 !bash .../oneclick_start.sh (实测 51.7 tok/s, 10.1GB, 草稿同样自动关)
#   极限长上下文: QUANT=Q4_K_M CTX=786432 !bash .../oneclick_start.sh (14.3GB)
#   强制开草稿(不建议 4bit): SPEC_TYPE=draft-simple !bash .../oneclick_start.sh
#   （T4 上 1M 上下文物理放不下: 1M KV@q4_0 也需 13GB+，加权重必爆）
#   （草稿模式下 CTX>300k 会 OOM: 草稿 KV/计算缓冲继承主上下文）
#   （思考模式默认关: aitun 隧道非流式请求 120s 硬超时, 开思考会 504; THINKING=on 恢复）
# ============================================================================
set -euo pipefail

BASE_DIR="/content/drive/MyDrive/spark-t4"
MODEL_DIR="/content/spark-t4-models"     # 模型永不进 Drive（会话重置后自动重下）
export BASE_DIR MODEL_DIR
export QUANT="${QUANT:-Q8_0}"             # 默认 Q8_0: 质量最佳+草稿+41%≈52tok/s(11.8GB); 4bit: QUANT=Q4_K_M(6.6GB,自动关草稿)
export CTX="${CTX:-262144}"                # 256k: 草稿投机解码最优档; 无草稿可设 500000
export PORT="${PORT:-8100}"              # 8080 被 Colab 环境占用
export NP="${NP:-1}"
export KVQ="${KVQ:-q4_0}"
export BATCH="${BATCH:-512}"
export UBATCH="${UBATCH:-512}"
export EXTRA_ARGS="${EXTRA_ARGS:--np 1 -b 512 -ub 512 -fa on --alias spark-x2.5-4b}"
export SPEC_TYPE_DEFAULT=1
# 草稿策略(按量化档自动): Q8_0 挂草稿(接受率0.89,+41%); 4bit 裸跑(接受率0.37-0.61,草稿反而慢一半)
if [[ "$QUANT" == Q8* ]]; then
    export SPEC_TYPE="${SPEC_TYPE:-draft-simple}"
else
    export SPEC_TYPE="${SPEC_TYPE:-none}"
fi
export DRAFT_MODEL="${DRAFT_MODEL:-/content/spark-t4-models/Spark-X2.5-1.7B-Q8_0.gguf}"
export DRAFT_NMAX="${DRAFT_NMAX:-8}"
export THINKING="${THINKING:-off}"       # 思考默认关: aitun 隧道非流式 120s 硬超时→504防护; THINKING=on 恢复
export BUILD_JOBS="${BUILD_JOBS:-2}"
TUNNEL="${TUNNEL:-1}"                    # 1=自动开公网隧道

GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; CYAN='\033[0;36m'; NC='\033[0m'
info() { echo -e "${GREEN}[一键启动]${NC} $*"; }
warn() { echo -e "${YELLOW}[等待]${NC} $*"; }
fail() { echo -e "${RED}[失败]${NC} $*"; exit 1; }

[[ -d /content/drive/MyDrive ]] || fail "Drive 未挂载，请先在 Colab 挂载 Google Drive"

# 1. 编译产物检查（在 Drive，跨会话持久）
BIN="$BASE_DIR/llama.cpp-spark/build/bin/llama-server"
if [[ -x "$BIN" ]]; then
    info "检测到已编译的 llama-server，跳过编译"
else
    warn "未检测到编译产物，开始完整编译（首次约 40~60 分钟）..."
    bash "$BASE_DIR/deploy_spark_t4.sh" build
fi

# 2. 模型检查（在本地盘，会话重置后自动重下，约 1~3 分钟）
if [[ -f "$MODEL_DIR/Spark-X2.5-4B-${QUANT}.gguf" ]]; then
    info "检测到模型: ${QUANT}"
else
    warn "模型不在本地盘（新会话正常现象），开始下载 ${QUANT} ..."
    mkdir -p "$MODEL_DIR"
    bash "$BASE_DIR/deploy_spark_t4.sh" download
fi

# 2.5 草稿模型检查（投机解码默认开启；缺失自动从 HF 公开仓补下）
if [[ "$SPEC_TYPE" != "none" && -n "$DRAFT_MODEL" && ! -f "$DRAFT_MODEL" ]]; then
    warn "草稿模型不在本地盘，从 HF 公开仓补下 Spark-X2.5-1.7B-Q8_0 (1.82GB)..."
    command -v hf >/dev/null 2>&1 || pip install -q -U huggingface_hub >/dev/null 2>&1 || true
    mkdir -p "$MODEL_DIR"
    hf download gasschina/Spark-X2.5-4B-build-cpp Spark-X2.5-1.7B-Q8_0.gguf \
        --local-dir "$MODEL_DIR" >/dev/null 2>&1 || true
fi
if [[ "$SPEC_TYPE" != "none" && -n "$DRAFT_MODEL" && -f "$DRAFT_MODEL" ]]; then
    info "草稿模型就绪: $DRAFT_MODEL（投机解码实测 +41%）"
elif [[ "$SPEC_TYPE" != "none" && -n "$DRAFT_MODEL" ]]; then
    warn "草稿模型下载失败，本次无草稿启动（稍后手动补下后重启即可）"
    export DRAFT_MODEL=""
fi

# 3. 起 API 服务（等待加载完成）
info "启动 API (ctx=$CTX port=$PORT quant=$QUANT kv=$KVQ)..."
bash "$BASE_DIR/spark_api.sh" start

# 4. 公网隧道
if [[ "$TUNNEL" == "1" ]]; then
    PORT="$PORT" bash "$BASE_DIR/tunnel_api.sh" start
fi

# 5. 快速自检
KEY="$(cat "$BASE_DIR/.api_key" 2>/dev/null || echo "")"
AUTH=(); [[ -n "$KEY" ]] && AUTH=(-H "Authorization: Bearer $KEY")
if curl -s --max-time 60 "http://127.0.0.1:$PORT/v1/chat/completions" \
    "${AUTH[@]}" -H "Content-Type: application/json" \
    -d '{"model":"spark","messages":[{"role":"user","content":"回复:OK"}],"max_tokens":64,"chat_template_kwargs":{"enable_thinking":false}}' \
    | grep -q '"content"'; then
    echo -e "${CYAN}════════════════════════════════════════════════${NC}"
    echo -e "${GREEN}  Spark-X2.5-4B 服务就绪 (E2E 自检通过)${NC}"
    echo -e "  本机:  http://127.0.0.1:$PORT/v1/chat/completions"
    [[ -n "$KEY" ]] && echo -e "  Key:   $KEY"
    echo -e "  上下文: $CTX tokens | 权重: $QUANT | KV: ${KVQ:-f16}"
    [[ -n "$DRAFT_MODEL" && "$SPEC_TYPE" != "none" ]] && echo -e "  投机解码: 1.7B 草稿 draft-simple (Q8_0 实测 +41% ≈52 tok/s, 接受率 0.89)"
    [[ "$SPEC_TYPE" == "none" ]] && echo -e "  投机解码: 关闭(4bit 裸跑 48.75 tok/s; 草稿仅对 Q8_0 有正收益)"
    echo -e "  思考模式: $([[ "$THINKING" == "off" ]] && echo '默认关闭(网关504防护, 单请求可传 enable_thinking:true)' || echo '默认开启')"
    [[ "$TUNNEL" == "1" ]] && echo -e "  公网:  https://aitun.cc/$(grep -oE 'aitun\.cc/[A-Za-z0-9]+' /content/aitun-api-tunnel.log 2>/dev/null | tail -1 | cut -d/ -f2)/v1/chat/completions"
    echo -e "  流式: 请求体加 \"stream\": true"
    echo -e "  开思考(单请求): 请求体加 \"chat_template_kwargs\":{\"enable_thinking\":true}"
    echo -e "${CYAN}════════════════════════════════════════════════${NC}"
else
    fail "自检未通过，查看: tail -30 $BASE_DIR/api-server.log"
fi