#!/usr/bin/env bash
# ============================================================================
# deploy_spark_t4.sh — Spark-X2.5-4B 在 NVIDIA T4 (16GB) 上的一键部署脚本
# 推理引擎: llama.cpp CUDA (XHToken 官方 fork)
#
# 【为什么必须编译，不能用现成的 llama.cpp / Ollama 二进制？】
#   1. spark2_5 架构（1 全注意力 + 3 滑窗 SWA 混合）尚未合并进 mainline
#      llama.cpp，直接用主分支或 Ollama 自带运行时会报
#      "unknown architecture" 无法加载 GGUF —— 官方模型卡指定要用
#      github.com/XHToken/llama.cpp 这个 fork。
#   2. Linux 默认编译产物不带 CUDA 后端，必须 -DGGML_CUDA=ON 重编。
#   这两点与 T4 无关，任何显卡都一样。T4 (SM 7.5 Turing) 是 llama.cpp
#   CUDA 后端完整支持的架构，编译一次即可，之后日常使用无需再编译。
#
# 【量化档位怎么选（T4 16GB 显存）】
#   Q8_0  (8bit)  4.07GB  质量近无损，T4 跑 32k~64k 上下文很从容
#   Q6_K  (6bit)  3.15GB  质量与速度折中
#   Q4_K_M(4bit)  2.42GB  速度最快、可冲 128k+ 长上下文  <- 默认
#   注意: 官方仓 XHToken/Spark-X2.5-4B-GGUF 里的 7.66GB 文件是 F16
#   全精度（不是 8bit），T4 上没必要用；8bit 用本脚本的默认 imatrix 仓。
#
# 【用法】
#   chmod +x deploy_spark_t4.sh
#   ./deploy_spark_t4.sh deps       # 1. 装系统依赖 + 构建工具
#   ./deploy_spark_t4.sh build      # 2. 克隆 fork 并编译 CUDA 版（约10~20分钟）
#   ./deploy_spark_t4.sh download   # 3. 下载量化模型（默认 Q4_K_M）
#   ./deploy_spark_t4.sh serve      # 4. 启动 OpenAI 兼容 API 服务
#   ./deploy_spark_t4.sh test       # 5. 发一条测试请求
#   ./deploy_spark_t4.sh status     # 查看运行状态
#   ./deploy_spark_t4.sh stop       # 停止服务
#   ./deploy_spark_t4.sh all        # 1~3 连跑（不含 serve）
#
#   换量化档位:  QUANT=Q4_K_M ./deploy_spark_t4.sh download
#   换上下文长度: CTX=65536 ./deploy_spark_t4.sh serve
#   无GPU机器只验证编译能否通过: GGML_CUDA=OFF ./deploy_spark_t4.sh build
# ============================================================================

set -euo pipefail

# ---------------------------- 可调配置 --------------------------------------
QUANT="${QUANT:-Q4_K_M}"                                  # Q4_K_M(默认) / Q8_0 / Q6_K / Q5_K_M ...
GGUF_REPO="${GGUF_REPO:-miifanboy/Spark-X2.5-4B-i1-GGUF}" # 含全档位 imatrix 量化
LLAMA_REPO="${LLAMA_REPO:-https://github.com/XHToken/llama.cpp.git}"
BASE_DIR="${BASE_DIR:-$HOME/spark-t4}"
CTX="${CTX:-32768}"                                       # 上下文长度, T4+Q8_0 建议 32k 起步
PORT="${PORT:-8080}"
NGL="${NGL:-999}"                                         # 全部层放 GPU
GGML_CUDA="${GGML_CUDA:-ON}"                              # 有 nvcc+GPU 用默认 ON；纯CPU验证编译改 OFF
BUILD_JOBS="${BUILD_JOBS:-$(nproc)}"                       # 内存紧张(<8GB)时建议 BUILD_JOBS=1
EXTRA_LLAMA_ARGS="${EXTRA_LLAMA_ARGS:-}"                  # 例如 "-ctk q8_0 -ctv q8_0"
HF_ENDPOINT_OPT="${HF_ENDPOINT_OPT:-}"                    # 国内镜像: export HF_ENDPOINT_OPT=https://hf-mirror.com

BUILD_DIR="$BASE_DIR/llama.cpp-spark"
MODEL_DIR="${MODEL_DIR:-$BASE_DIR/models}"                # 模型目录可独立于 BASE_DIR（如 Colab 上放本地盘省 Drive 配额）
LOG_FILE="$BASE_DIR/llama-server.log"
PID_FILE="$BASE_DIR/llama-server.pid"
MODEL_FILE="$MODEL_DIR/Spark-X2.5-4B-${QUANT}.gguf"

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

preflight() {
    if ! command -v nvidia-smi >/dev/null 2>&1; then
        if [[ "${GGML_CUDA^^}" == "ON" ]]; then
            fail "未检测到 nvidia-smi，请先装好 NVIDIA 驱动（GCP 可用 Deep Learning VM 镜像自带驱动）"
        else
            warn "未检测到 nvidia-smi（GGML_CUDA=OFF 纯 CPU 验证模式，继续）"
        fi
        return 0
    fi
    nvidia-smi --query-gpu=name,memory.total --format=csv,noheader | grep -qi "t4" \
        || warn "未检测到 T4（当前 GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)），脚本仍可继续"
}

# ---------------------------- 1. 系统依赖 ------------------------------------
do_deps() {
    preflight
    info "安装系统构建依赖（需要 sudo）..."
    sudo apt-get update -y
    sudo apt-get install -y build-essential cmake git curl python3 python3-pip python3-venv

    # 编译 CUDA 后端需要 nvcc；GCP DLVM 镜像一般自带
    if ! command -v nvcc >/dev/null 2>&1; then
        warn "未检测到 nvcc (CUDA Toolkit)"
        read -r -p "是否用 apt 安装 nvidia-cuda-toolkit（Ubuntu 源，版本较旧但足够编 T4）？[y/N] " ans
        if [[ "${ans,,}" == "y" ]]; then
            sudo apt-get install -y nvidia-cuda-toolkit
        else
            fail "请自行安装 CUDA Toolkit 12.x（https://developer.nvidia.com/cuda-downloads）后重试"
        fi
    fi
    nvcc --version | tail -1

    info "安装 Hugging Face 下载工具..."
    python3 -m pip install -q -U "huggingface_hub[cli]" || \
        python3 -m pip install -q --user -U "huggingface_hub[cli]"
    info "依赖安装完成。"
}

# ---------------------------- 2. 编译 fork -----------------------------------
do_build() {
    preflight
    if [[ "${GGML_CUDA^^}" == "ON" ]]; then
        command -v nvcc >/dev/null 2>&1 || fail "CUDA=ON 需要 nvcc，请先执行 ./deploy_spark_t4.sh deps；纯 CPU 验证用 GGML_CUDA=OFF ./deploy_spark_t4.sh build"
    fi
    mkdir -p "$BASE_DIR"

    if [[ -d "$BUILD_DIR/.git" ]]; then
        info "已存在 fork 目录，拉取最新代码..."
        git -C "$BUILD_DIR" pull --ff-only || warn "git pull 失败，使用现有代码继续"
    else
        info "克隆官方 llama.cpp fork (XHToken/llama.cpp)..."
        git clone --depth 1 "$LLAMA_REPO" "$BUILD_DIR"
    fi

    if [[ "${GGML_CUDA^^}" == "ON" ]]; then
        info "配置 CMake (CUDA 后端, 针对 T4/SM75)..."
        cmake -S "$BUILD_DIR" -B "$BUILD_DIR/build" \
            -DGGML_CUDA=ON \
            -DCMAKE_BUILD_TYPE=Release
    else
        info "配置 CMake (纯 CPU 验证模式, 不编译 CUDA 内核)..."
        cmake -S "$BUILD_DIR" -B "$BUILD_DIR/build" \
            -DGGML_CUDA=OFF \
            -DCMAKE_BUILD_TYPE=Release
    fi

    info "开始编译，线程数: $BUILD_JOBS，约 10~20 分钟..."
    cmake --build "$BUILD_DIR/build" --config Release -j "$BUILD_JOBS"

    local bin="$BUILD_DIR/build/bin/llama-server"
    [[ -x "$bin" ]] || fail "编译完成但未找到 llama-server，请检查上方日志"
    info "编译成功: $bin"
}

# ---------------------------- 3. 下载模型 ------------------------------------
do_download() {
    mkdir -p "$MODEL_DIR"
    if [[ -f "$MODEL_FILE" ]]; then
        info "模型已存在: $MODEL_FILE（如需重下请先删除）"
        return 0
    fi
    [[ -n "$HF_ENDPOINT_OPT" ]] && export HF_ENDPOINT="$HF_ENDPOINT_OPT"

    local dl
    if command -v hf >/dev/null 2>&1; then dl="hf download"
    elif command -v huggingface-cli >/dev/null 2>&1; then dl="huggingface-cli download"
    else fail "缺少 huggingface_hub，请先执行 ./deploy_spark_t4.sh deps"; fi

    info "下载 $GGUF_REPO / Spark-X2.5-4B-${QUANT}.gguf ..."
    $dl "$GGUF_REPO" "Spark-X2.5-4B-${QUANT}.gguf" --local-dir "$MODEL_DIR"

    [[ -f "$MODEL_FILE" ]] || fail "下载后未找到 $MODEL_FILE"
    local size
    size=$(du -m "$MODEL_FILE" | cut -f1)
    info "下载完成: ${size} MB"
    # 档位合理性提示（T4 16GB）
    if [[ "$QUANT" == "Q8_0" ]]; then
        info "Q8_0 (8bit) 在 T4 上建议上下文 32k~64k；要冲 128k 请换 QUANT=Q4_K_M 重新 download"
    fi
}

# ---------------------------- 4. 启动服务 ------------------------------------
do_serve() {
    [[ -x "$BUILD_DIR/build/bin/llama-server" ]] || fail "请先执行 build"
    [[ -f "$MODEL_FILE" ]] || fail "请先执行 download"
    do_stop 2>/dev/null || true

    info "启动 llama-server: ctx=$CTX port=$PORT quant=$QUANT"
    nohup "$BUILD_DIR/build/bin/llama-server" \
        -m "$MODEL_FILE" \
        -ngl "$NGL" \
        -c "$CTX" \
        --jinja \
        --host 0.0.0.0 \
        --port "$PORT" \
        $EXTRA_LLAMA_ARGS \
        > "$LOG_FILE" 2>&1 &
    echo $! > "$PID_FILE"

    sleep 5
    if kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
        info "服务已启动 (PID $(cat "$PID_FILE"))，日志: $LOG_FILE"
        info "API 地址: http://<本机IP>:$PORT/v1/chat/completions"
        info "思考模式默认开启，响应会包含推理过程；max_tokens 建议给大一些"
        warn "若服务几秒后退出，执行 ./deploy_spark_t4.sh status 看日志"
        warn "常见原因: ①--jinja 参数不被该 fork 支持 → EXTRA_LLAMA_ARGS 去掉重试"
        warn "          ②显存不足 → 调小 CTX 或换 QUANT=Q4_K_M"
    else
        fail "启动失败，最近日志:\n$(tail -20 "$LOG_FILE")"
    fi
}

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() {
    nvidia-smi --query-gpu=name,memory.used,memory.total,utilization.gpu --format=csv,noheader || true
    if [[ -f "$PID_FILE" ]] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
        info "服务运行中 (PID $(cat "$PID_FILE"))"
        tail -5 "$LOG_FILE" 2>/dev/null || true
    else
        warn "服务未运行"
    fi
}

# ---------------------------- 5. 测试请求 ------------------------------------
do_test() {
    info "发送测试请求到 http://127.0.0.1:$PORT ..."
    curl -s "http://127.0.0.1:$PORT/v1/chat/completions" \
        -H "Content-Type: application/json" \
        -d '{
            "model": "spark",
            "messages": [{"role": "user", "content": "用一句话介绍你自己"}],
            "max_tokens": 1024,
            "temperature": 1.0,
            "top_p": 0.95
        }' | python3 -c "
import json, sys
try:
    d = json.load(sys.stdin)
    print(d['choices'][0]['message'].get('reasoning_content') or '')
    print(d['choices'][0]['message']['content'])
except Exception as e:
    print('解析失败:', e); sys.exit(1)
"
    info "官方推荐采样参数: temperature=1.0, top_p=0.95, top_k=-1 (llama.cpp 中 top_k=0 表示禁用)"
}

# ---------------------------- 入口 ------------------------------------------
case "${1:-help}" in
    deps)     do_deps ;;
    build)    do_build ;;
    download) do_download ;;
    serve)    do_serve ;;
    stop)     do_stop ;;
    status)   do_status ;;
    test)     do_test ;;
    all)      do_deps; do_build; do_download
              info "全部就绪！执行 ./deploy_spark_t4.sh serve 启动服务" ;;
    *)        sed -n '2,40p' "$0" | sed 's/^# \{0,1\}//' ;;
esac