#!/usr/bin/env bash
# ============================================================================
# NVIDIA A30 24GB 一键部署脚本：MiniCPM-V 2.6 8B + SGLang 推理 + LLaMA-Factory 训练
#
# 适用环境:
#   · Ubuntu 22.04 (kernel 5.15+)
#   · Python 3.10+ (自带 venv)
#   · NVIDIA Driver >= 550 (CUDA 12.4)
#   · GPU: A30 24GB (Ampere SM80, BF16 + FP8 + AWQ INT4 全支持)
#
# 用法:
#   chmod +x a30_setup.sh
#
#   # 第一步：环境检查 (不安装, 只检查)
#   ./a30_setup.sh check
#
#   # 第二步：创建 venv + 装依赖 (3-5 分钟)
#   ./a30_setup.sh install
#
#   # 第三步：下载 MiniCPM-V 2.6 8B 模型 (ModelScope 国内源, 10 分钟)
#   ./a30_setup.sh download-model
#
#   # 第四步：生成启动脚本 + 启动推理服务 (显存 ~9GB, 12 并发)
#   ./a30_setup.sh serve
#
#   # 一步到位 (check+install+download+serve, 15-25 分钟)
#   ./a30_setup.sh all
#
#   # 生成 systemd 服务 (开机自启, 守护重启)
#   ./a30_setup.sh install-systemd
#
# ============================================================================
set -euo pipefail

# ---------- 配置 ----------
BASE_DIR="${BASE_DIR:-/data/volc_minicpm}"          # 安装根目录
VENV_DIR="${BASE_DIR}/venv"                          # 虚拟环境
MODEL_DIR="${BASE_DIR}/models"                       # 模型保存目录
CKPT_DIR="${BASE_DIR}/checkpoints"                   # LoRA checkpoint
LOG_DIR="${BASE_DIR}/logs"                           # 日志
PORT="${PORT:-30000}"                                # SGLang 端口
HOST="${HOST:-0.0.0.0}"
MODEL_ID="${MODEL_ID:-openbmb/MiniCPM-V-2_6}"        # 8B 旗舰版
MODEL_NAME="${MODEL_NAME:-minicpmv}"
QUANT="${QUANT:-awq}"                                # A30 默认 AWQ INT4 量化
MEM_FRAC="${MEM_FRAC:-0.90}"
DEPLOY_SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"   # a30_deploy.py 所在目录

RED="\033[91m"; GREEN="\033[92m"; YELLOW="\033[93m"; CYAN="\033[96m"; NC="\033[0m"
log()    { echo -e "${GREEN}[INFO]${NC}  $*"; }
warn()   { echo -e "${YELLOW}[WARN]${NC}  $*"; }
err()    { echo -e "${RED}[ERR ]${NC}  $*" 1>&2; }
header() { echo -e "\n${CYAN}========== $* ==========${NC}\n"; }

# ---------- 步骤: 环境检查 ----------
step_check() {
    header "环境检查"

    # OS 检查
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        log "OS:        $PRETTY_NAME"
    fi
    log "Kernel:    $(uname -r)"
    log "Arch:      $(uname -m)"

    # Python
    PY3="$(command -v python3 || true)"
    if [[ -z "$PY3" ]]; then
        err "未找到 python3, 请先 apt install python3 python3-venv python3-pip"
        exit 1
    fi
    PY_VER="$($PY3 -c 'import sys; print("%d.%d"%(sys.version_info.major,sys.version_info.minor))')"
    log "Python:    $PY3 $PY_VER"
    if [[ ! "$PY_VER" =~ ^(3\.10|3\.11|3\.12)$ ]]; then
        warn "推荐 Python 3.10/3.11/3.12, 当前版本可能有兼容性问题"
    fi
    if ! $PY3 -c "import venv" 2>/dev/null; then
        err "缺少 venv 模块, 请执行: apt install -y python3-venv"
        exit 1
    fi

    # NVIDIA Driver / CUDA
    if ! command -v nvidia-smi >/dev/null 2>&1; then
        err "未找到 nvidia-smi, NVIDIA 驱动未安装?"
        exit 1
    fi
    DRV_VER="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)"
    CUDA_VER="$(nvidia-smi | head -3 | grep -oE "CUDA Version:[[:space:]]*[0-9.]+" | awk '{print $3}')"
    GPU_NAME="$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)"
    GPU_MEM="$(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -1 | tr -d ' ')"
    log "GPU:       $GPU_NAME ($GPU_MEM)"
    log "Driver:    $DRV_VER"
    log "CUDA:      $CUDA_VER (驱动支持上限)"

    # 显存 >= 24GB 推荐 AWQ INT4 8B, 否则降级
    MEM_GB="$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1 | awk '{printf "%.0f", $1/1024}')"
    if (( MEM_GB >= 24 )); then
        log "✅ 显存充足 (${MEM_GB}GB), 推荐 MiniCPM-V 8B + AWQ INT4 (9GB) + 12 并发"
    elif (( MEM_GB >= 16 )); then
        warn "显存 ${MEM_GB}GB, 推荐 MiniCPM-V 4B BF16 (16GB) 或 8B INT4 (9GB)"
    else
        warn "显存 ${MEM_GB}GB, 只能 MiniCPM-V 2B (3GB INT4)"
    fi

    # 基本系统依赖
    for pkg in git curl ca-certificates; do
        if ! command -v $pkg >/dev/null 2>&1; then
            warn "缺少 $pkg, 建议 apt install -y $pkg"
        fi
    done

    # 磁盘空间 (至少 80GB)
    FREE_GB="$(df -BG --output=avail "${BASE_DIR%/*}" 2>/dev/null | tail -1 | tr -d ' G' || echo 0)"
    log "磁盘可用:  ${FREE_GB} GB (至少 80GB 用于存放 8B BF16 模型+权重缓存)"
    if (( FREE_GB < 80 )); then
        warn "磁盘空间可能不足, 建议清理后再继续"
    fi

    echo ""
    log "✅ 环境检查完成"
}

# ---------- 步骤: 安装依赖 ----------
step_install() {
    header "创建目录 & 虚拟环境 + 安装依赖"

    mkdir -p "$BASE_DIR" "$VENV_DIR" "$MODEL_DIR" "$CKPT_DIR" "$LOG_DIR"

    if [[ ! -f "$VENV_DIR/bin/activate" ]]; then
        log "创建 venv: $VENV_DIR"
        python3 -m venv "$VENV_DIR"
    fi
    # shellcheck disable=SC1091
    source "$VENV_DIR/bin/activate"

    log "升级 pip/setuptools/wheel"
    $VENV_DIR/bin/pip install --upgrade pip setuptools wheel 2>&1 | tail -5

    # 1) 关键: torch + CUDA 12.1 (NVIDIA Driver 12.4 向下兼容 12.1)
    log "安装 PyTorch (CUDA 12.1) ..."
    $VENV_DIR/bin/pip install --no-cache-dir \
        --extra-index-url https://download.pytorch.org/whl/cu121 \
        torch==2.3.1 torchvision==0.18.1 2>&1 | tail -5

    # 验证 torch CUDA
    $VENV_DIR/bin/python - <<PY
import torch
print(f"torch.version = {torch.__version__}")
print(f"torch.cuda.is_available = {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"torch.cuda.device_count = {torch.cuda.device_count()}")
    print(f"torch.backends.cudnn.enabled = {torch.backends.cudnn.enabled}")
    print(f"torch.version.cuda = {torch.version.cuda}")
    gbn = torch.cuda.get_device_properties(0).total_memory/1024**3
    print(f"GPU 显存 = {gbn:.1f} GB")
    assert gbn >= 20, f"显存不足 {gbn:.1f} GB < 20GB, A30 应该是 24GB"
    print("✅ torch + CUDA OK")
PY

    # 2) SGLang 推理框架
    log "安装 SGLang (推理框架) ..."
    $VENV_DIR/bin/pip install --no-cache-dir "sglang[all]" 2>&1 | tail -10

    # 3) LLaMA-Factory 训练框架
    log "安装 LLaMA-Factory (训练) + ModelScope ..."
    $VENV_DIR/bin/pip install --no-cache-dir \
        "llama-factory[modelscope]" modelscope transformers peft accelerate datasets 2>&1 | tail -10

    # 4) 其他依赖
    log "安装其他常用依赖 (httpx / pillow / opencv / numpy ...)"
    $VENV_DIR/bin/pip install --no-cache-dir \
        httpx pillow opencv-python-headless numpy pyyaml jinja2 uvicorn fastapi 2>&1 | tail -10

    # 验证关键版本
    echo ""
    log "📦 关键包版本:"
    for pkg in torch sglang llama-factory modelscope transformers peft accelerate; do
        v="$($VENV_DIR/bin/pip show "$pkg" 2>/dev/null | grep -i '^Version' | awk '{print $2}' || echo '(未安装)')"
        printf "    %-22s %s\n" "$pkg" "$v"
    done

    echo ""
    log "✅ 依赖安装完成, 虚拟环境: $VENV_DIR"
    log "💡 激活命令: source $VENV_DIR/bin/activate"
}

# ---------- 步骤: 下载模型 ----------
step_download_model() {
    header "下载 MiniCPM-V 2.6 8B (ModelScope 国内加速)"

    if [[ ! -f "$VENV_DIR/bin/activate" ]]; then
        err "先执行 ./a30_setup.sh install 安装依赖"
        exit 1
    fi
    source "$VENV_DIR/bin/activate"

    mkdir -p "$MODEL_DIR"

    $VENV_DIR/bin/python - <<PY
import sys, os
from modelscope import snapshot_download

dst = os.environ["MODEL_DIR"]
mid = os.environ["MODEL_ID"]
print(f"📥 下载: {mid}")
print(f"   保存目录: {dst}")
try:
    local = snapshot_download(
        mid,
        cache_dir=dst + "/.ms_cache",
        local_dir=os.path.join(dst, mid.split("/")[-1]),
    )
    print(f"\n✅ 模型已下载: {local}")
except Exception as e:
    print(f"❌ 下载失败: {e}", file=sys.stderr)
    sys.exit(1)
PY

    echo ""
    log "模型文件大小:"
    du -sh "$MODEL_DIR"/*/ 2>/dev/null | head
}

# ---------- 步骤: 生成启动脚本 + 启动 ----------
step_serve() {
    header "生成 SGLang 启动脚本 & 启动推理服务 (端口=$PORT)"

    if [[ ! -f "$VENV_DIR/bin/activate" ]]; then
        err "先执行 ./a30_setup.sh install 安装依赖"
        exit 1
    fi
    source "$VENV_DIR/bin/activate"

    # 找到模型实际目录 (如果存在多个, 选最大的)
    MODEL_ACTUAL="$($VENV_DIR/bin/python - <<PY
import os, glob
d = os.environ["MODEL_DIR"]
cands = [x for x in glob.glob(os.path.join(d,"*")) if os.path.isdir(x) and ".ms_cache" not in x]
if len(cands) == 0:
    print("", end="")
    exit(0)
def size(p): return sum(os.path.getsize(os.path.join(r,f)) for r,_,fs in os.walk(p) for f in fs)
cands.sort(key=size, reverse=True)
print(cands[0])
PY
)"
    if [[ -z "$MODEL_ACTUAL" ]]; then
        warn "没有在 $MODEL_DIR 下找到模型, 先执行 ./a30_setup.sh download-model"
        warn "(或者手动指定 MODEL_ACTUAL=/path 后再执行)"
        exit 1
    fi
    log "实际模型路径: $MODEL_ACTUAL"

    # 生成启动脚本 (用 a30_deploy.py 生成)
    if [[ -f "$DEPLOY_SCRIPT_DIR/a30_deploy.py" ]]; then
        log "用 a30_deploy.py 生成启动命令 + 一键脚本"
        $VENV_DIR/bin/python "$DEPLOY_SCRIPT_DIR/a30_deploy.py" serve-sglang \
            --model-path "$MODEL_ACTUAL" \
            --model-name "$MODEL_NAME" \
            --port "$PORT" \
            --host "$HOST" \
            --quant "$QUANT" \
            --mem-fraction "$MEM_FRAC" \
            --launch-script "$BASE_DIR/start_sglang.sh"
    else
        warn "找不到 a30_deploy.py, 将手写启动脚本"
        LAUNCHER="$BASE_DIR/start_sglang.sh"
        cat >"$LAUNCHER" <<BASH
#!/usr/bin/env bash
set -e
source "$VENV_DIR/bin/activate"
cd "$BASE_DIR"
exec python3 -m sglang.launch_server \
    --model-path "$MODEL_ACTUAL" \
    --trust-remote-code \
    --served-model-name "$MODEL_NAME" \
    --port "$PORT" \
    --host "$HOST" \
    --mem-fraction-static "$MEM_FRAC" \
    --max-total-tokens 8192 \
    --quantization awq \
    --dp-size 1 \
    >> "$LOG_DIR/sglang.log" 2>&1
BASH
        chmod +x "$LAUNCHER"
        log "已生成启动脚本: $LAUNCHER"
    fi

    # 启动服务 (后台)
    log "🚀 启动 SGLang 推理服务 (nohup 后台, 日志=$LOG_DIR/sglang.log) ..."
    LAUNCHER="$BASE_DIR/start_sglang.sh"
    chmod +x "$LAUNCHER"
    nohup "$LAUNCHER" >> "$LOG_DIR/sglang.log" 2>&1 &
    SERV_PID=$!
    echo "$SERV_PID" > "$BASE_DIR/sglang.pid"
    log "服务 PID=$SERV_PID"

    # 等服务起来 (最多 180 秒, 加载大模型慢)
    log "⏳ 等待服务就绪 (加载 8B AWQ INT4 需要 30-90 秒) ..."
    for i in {1..36}; do
        sleep 5
        if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 || \
           curl -sf "http://127.0.0.1:$PORT/v1/models" >/dev/null 2>&1; then
            log "✅ 服务已就绪! PID=$SERV_PID"
            break
        fi
        # 进程还活着吗
        if ! kill -0 "$SERV_PID" 2>/dev/null; then
            err "服务进程已退出, 查看日志: tail -30 $LOG_DIR/sglang.log"
            exit 1
        fi
        echo -n "."
    done

    echo ""
    echo "========================================="
    echo "✅ SGLang 推理服务部署完成"
    echo "   · 进程 PID:  $(cat "$BASE_DIR/sglang.pid")"
    echo "   · API Base:  http://<A30_IP>:$PORT/v1"
    echo "   · 接口文档:  http://<A30_IP>:$PORT/docs"
    echo "   · 启动脚本:  $BASE_DIR/start_sglang.sh"
    echo "   · 停止服务:  kill \$(cat $BASE_DIR/sglang.pid)"
    echo "   · 实时日志:  tail -f $LOG_DIR/sglang.log"
    echo "========================================="
}

# ---------- 步骤: 生成 systemd 服务 ----------
step_install_systemd() {
    header "生成 systemd 服务 (开机自启 + 守护重启)"

    LAUNCHER="$BASE_DIR/start_sglang.sh"
    if [[ ! -x "$LAUNCHER" ]]; then
        err "先执行 ./a30_setup.sh serve 生成启动脚本"
        exit 1
    fi

    SVC_FILE="/etc/systemd/system/minicpm-sglang.service"
    cat >"$SVC_FILE" <<SVC
[Unit]
Description=MiniCPM-V SGLang Inference Service (A30)
After=network.target nvidia-persistenced.service

[Service]
Type=simple
User=root
Group=root
WorkingDirectory=$BASE_DIR
ExecStart=$LAUNCHER
Restart=always
RestartSec=5
KillSignal=SIGTERM
TimeoutStopSec=30
StandardOutput=append:$LOG_DIR/sglang.log
StandardError=append:$LOG_DIR/sglang.log
LimitMEMLOCK=infinity
LimitNOFILE=65536

[Install]
WantedBy=multi-user.target
SVC
    log "✅ 生成 service 文件: $SVC_FILE"

    systemctl daemon-reload
    systemctl enable minicpm-sglang.service
    log "✅ 已设置开机自启 (minicpm-sglang.service)"
    log ""
    log "手动控制命令:"
    log "   systemctl start    minicpm-sglang   # 启动"
    log "   systemctl stop     minicpm-sglang   # 停止"
    log "   systemctl restart  minicpm-sglang   # 重启"
    log "   systemctl status   minicpm-sglang   # 状态"
    log "   journalctl -u minicpm-sglang -f     # 实时日志"
}

# ---------- 主入口 ----------
cmd="${1:-all}"
case "$cmd" in
    check)           step_check ;;
    install)         step_check && step_install ;;
    download-model)  step_download_model ;;
    serve)           step_serve ;;
    install-systemd) step_install_systemd ;;
    all)             step_check && step_install && step_download_model && step_serve ;;
    -h|--help|help)
        sed -n '2,30p' "$0" ;;
    *)
        err "未知命令: $cmd, 可用: check | install | download-model | serve | install-systemd | all"
        exit 1 ;;
esac