#!/usr/bin/env bash
set -u
# ============================================================
#  vLLM 生产级部署 一键脚本（macOS / Linux / WSL）
#  公开源码，欢迎审查 —— 不放心可先复制给 AI 判断
#  前置：NVIDIA GPU + CUDA（无 GPU 也可装，CPU 推理较慢）
#  国内自动开启 ModelScope 模型源
# ============================================================

GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; CYAN='\033[0;36m'; NC='\033[0m'
info()  { echo -e "${GREEN}[INFO]${NC} $1"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; }
step()  { echo -e "${CYAN}[STEP]${NC} $1"; }

detect_network() {
  step "检测网络环境（国内 / 国外）..."
  if curl -fsI --max-time 5 "https://claude.ai" >/dev/null 2>&1; then
    info "可直连 claude.ai，判定为海外网络"
    echo "overseas"
  else
    warn "无法直连 claude.ai，判定为国内网络环境（将自动切换国内镜像 + ModelScope 模型源）"
    echo "domestic"
  fi
}

check_gpu() {
  step "检查 GPU 环境..."
  if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
    info "检测到 NVIDIA GPU："
    nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null | head -3
    return 0
  fi
  warn "未检测到 NVIDIA GPU（nvidia-smi 不可用）。vLLM 仍可安装，但推理走 CPU 会明显更慢。"
  echo "  · 有独显的笔记本/服务器请先装好 NVIDIA 驱动（命令 nvidia-smi 能输出即 OK）"
  echo "  · Windows 用户请在 WSL 中运行本脚本（NVIDIA 官方支持 WSL CUDA）"
  echo "  · 没有 GPU 也可以继续（CPU 模式），或考虑用云端 API 更快"
  return 1
}

check_python312() {
  step "检查 Python 3.12..."
  local found=""
  for p in python3.12 python3 python3.13; do
    if command -v "$p" >/dev/null 2>&1 && "$p" --version 2>&1 | grep -qE "3\.(12|13)"; then
      found="$p"; break
    fi
  done
  if [ -n "$found" ]; then
    PY="$found"; info "使用 $PY：$($PY --version 2>&1)"; return 0
  fi
  error "需要 Python 3.12+。请先安装（https://www.python.org/downloads/ 或 distro 包管理器）。"
  return 1
}

setup_and_install() {
  local net="$1"
  step "创建虚拟环境并安装 vLLM（依赖较大，可能耗时 5-15 分钟，可耐心等或先干别的）..."
  mkdir -p vllm-serve && cd vllm-serve || return 1
  $PY -m venv .venv 2>/dev/null || { error "创建虚拟环境失败"; return 1; }
  # shellcheck disable=SC1091
  source .venv/bin/activate || { warn "激活虚拟环境失败，继续使用系统环境"; }

  local pypi="https://pypi.org/simple"
  if [ "$net" = "domestic" ]; then
    pypi="https://pypi.tuna.tsinghua.edu.cn/simple"
  fi

  step "安装 vllm（清华镜像优先，失败回退官方源）..."
  if pip install "vllm" -i "$pypi" --timeout 120 >/dev/null 2>&1; then
    info "vLLM 安装成功"
    return 0
  fi
  warn "镜像源安装失败，改用官方源重试（依赖较多请耐心）..."
  if pip install "vllm" --timeout 120 >/dev/null 2>&1; then
    info "vLLM 安装成功（官方源）"
    return 0
  fi
  error "vLLM 安装失败。常见原因与处理："
  echo "  1. 内存不足：vLLM + PyTorch 安装需要 8GB+ 内存，请关掉其他程序"
  echo "  2. Python 版本不匹配：要求 3.12+（当前 $($PY --version 2>&1)）"
  echo "  3. 网络问题：手动执行 python3 -m pip install vllm -i https://pypi.tuna.tsinghua.edu.cn/simple"
  return 1
}

start_service() {
  step "生成启动脚本与用法说明（模型下载走 ModelScope 可选）..."
  cat > start.sh <<'SHEOF'
#!/usr/bin/env bash
# vLLM 启动示例（模型按需改）：
#   国内（ModelScope 源）：VLLM_USE_MODELSCOPE=True ./start.sh
#   海外（Hugging Face 源）：./start.sh
set -e
cd "$(dirname "$0")"
if [ -f .venv/bin/activate ]; then source .venv/bin/activate; fi
export VLLM_USE_MODELSCOPE=${VLLM_USE_MODELSCOPE:-False}
# 首个模型可换成任意 HF/ModelScope 模型名，如 Qwen/Qwen2.5-1.5B-Instruct、deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
vllm serve Qwen/Qwen2.5-1.5B-Instruct --port 8000
SHEOF
  chmod +x start.sh

  cat > test_client.py <<'PYEOF'
"""调用本地 vLLM 服务的 OpenAI 兼容客户端（和服务起好后另开终端运行）"""
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
resp = client.chat.completions.create(
    model="Qwen/Qwen2.5-1.5B-Instruct",
    messages=[{"role": "user", "content": "你好！用一句话介绍你自己。"}],
)
print(resp.choices[0].message.content)
PYEOF

  info "已生成 vllm-serve/start.sh 与 vllm-serve/test_client.py"
}

main() {
  echo "============================================"
  echo "  vLLM 生产级部署 一键脚本"
  echo "  适用于 macOS / Linux / WSL"
  echo "============================================"
  echo ""

  local NET
  NET=$(detect_network)

  if [ "$NET" = "domestic" ]; then
    export VLLM_USE_MODELSCOPE=True
    info "已启用 ModelScope 模型源（VLLM_USE_MODELSCOPE=True），国内下载模型快"
  fi

  check_gpu || warn "继续安装（CPU 模式可跑但慢；想省事可在文末考虑云端 API 路线）"
  if ! check_python312; then exit 1; fi
  if ! setup_and_install "$NET"; then
    error "环境准备失败，脚本终止。请按上方提示手动处理后重跑。"
    exit 1
  fi

  start_service

  echo ""
  echo "============================================"
  info "全部完成！"
  echo "  启动服务：  cd vllm-serve && VLLM_USE_MODELSCOPE=True ./start.sh"
  echo "  验证服务：  curl http://localhost:8000/v1/models"
  echo "  测试调用：  另开终端 cd vllm-serve && source .venv/bin/activate && python test_client.py"
  echo "  换模型：    编辑 start.sh 里 vllm serve 的模型名"
  echo ""
  echo "  想省 GPU/精力：云端 OpenAI 兼容 API（如 https://cloudzone-api.cyou/，90+ 模型按量计费）"
  echo "============================================"
}

main "$@"