#!/usr/bin/env bash
set -euo pipefail

############################################
# vLLM(Qwen-VL) + 批处理流水线（防截断·参数加强）
# - 上下文窗放大：--max-model-len 40960（Qwen3 默认）
# - 显存利用提升：--gpu-memory-utilization 0.95
# - 并发适度下降：--max-num-seqs 8（给长输出更稳定）
# - TP 自适配（整除 attention heads）
# - setsid 进程组 + 组杀，退出优雅清理
############################################

# ======== 可改参数 ========
CSV_PATH="${CSV_PATH:-/mnt/data_nas/guanfb/Edit/CoT_Edit/8BTest_CoT_box/data/Test.csv}"
OUTDIR="${OUTDIR:-./outputs}"
MODEL_SIZE="${MODEL_SIZE:-8b}"                 # 8b | 32b
API_PORT="${API_PORT:-8007}"                   # 8B默认8007，32B可用8032
RESULTS_CSV="${RESULTS_CSV:-$OUTDIR/results.csv}"
DRY_RUN="${DRY_RUN:-0}"
RESUME="${RESUME:-1}"
WAIT_SECS="${WAIT_SECS:-600}"

# === 客户端并发（传给 run.py） ===
CLIENT_HTTP_WORKERS="${CLIENT_HTTP_WORKERS:-2}"    # 建议 2~4
CLIENT_HTTP_TIMEOUT="${CLIENT_HTTP_TIMEOUT:-1200}" # 单请求超时
CLIENT_MAX_TOKENS="${CLIENT_MAX_TOKENS:-4096}"     # 生成上限（配合防截断）

# GPU 绑定（多卡逗号分隔）
GPU_IDS="${GPU_IDS:-0,1,2,3}"
export CUDA_VISIBLE_DEVICES="${GPU_IDS}"
IFS=',' read -ra _g <<< "$GPU_IDS"
NUM_GPUS="${#_g[@]}"

# vLLM 服务参数（防截断关键项）
GPU_UTIL="${GPU_UTIL:-0.95}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-8}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-40960}"

SERVED_NAME="${SERVED_NAME:-qwen_vl}"

# 模型权重
if [[ "$MODEL_SIZE" == "32b" ]]; then
  MODEL_NAME="${MODEL_NAME:-/path/to/Qwen3-VL-32B-Instruct}"
  API_PORT="${API_PORT:-8032}"
else
  MODEL_NAME="${MODEL_NAME:-/mnt/data_nas/guanfb/Edit/CoT_Edit/8BTest_CoT_box/Qwen3-VL-8B-Instruct}"
fi

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOG_DIR="$OUTDIR/logs"
mkdir -p "$LOG_DIR" "$OUTDIR"

echo "[1/6] Installing Python deps if needed..."
pip install -r "$ROOT_DIR/requirements.txt" >/dev/null || true

# ===== TP 自适配（读取 num_attention_heads，整除选 TP） =====
echo "[2/6] Probing model heads & selecting tensor parallel size..."
NUM_HEADS=$(python - <<'PY'
import json, os
m=os.environ.get("MODEL_NAME","")
cfg=os.path.join(m,"config.json")
h=None
try:
    with open(cfg,"r",encoding="utf-8") as f:
        j=json.load(f)
        h=j.get("num_attention_heads") or j.get("n_head") or j.get("num_key_value_heads")
except Exception:
    pass
if h is None:
    h = 64 if any(s in m.lower() for s in ("32","32b")) else 32
print(h)
PY
)
pick_tp() { local heads="$1"; local gpus="$2"; for ((d=gpus; d>=1; d--)); do (( heads % d == 0 )) && { echo "$d"; return; }; done; echo 1; }
TP_SIZE="${TP_SIZE:-$(pick_tp "$NUM_HEADS" "$NUM_GPUS")}"
echo "   - CUDA_VISIBLE_DEVICES=${GPU_IDS} (num_gpus=${NUM_GPUS})"
echo "   - num_attention_heads=${NUM_HEADS}  => tensor_parallel_size=${TP_SIZE}"

API_BASE="http://127.0.0.1:${API_PORT}/v1"
echo "[3/6] Starting vLLM server: $MODEL_NAME on port $API_PORT (served as '$SERVED_NAME')"

# 端口占用兜底清理
if lsof -iTCP:${API_PORT} -sTCP:LISTEN -Pn >/dev/null 2>&1; then
  echo "[CLEANUP] Port ${API_PORT} is in use. Killing existing server..."
  old_pids=$(lsof -iTCP:${API_PORT} -sTCP:LISTEN -t || true)
  if [[ -n "$old_pids" ]]; then
    kill $old_pids 2>/dev/null || true
    sleep 2
    for p in $old_pids; do ps -p $p >/dev/null 2>&1 && kill -9 $p 2>/dev/null || true; done
  fi
fi

VLLM_PID_FILE="$LOG_DIR/vllm_${MODEL_SIZE}.pid"
VLLM_PGID_FILE="$LOG_DIR/vllm_${MODEL_SIZE}.pgid"

# setsid -> 独立进程组，方便组杀
setsid bash -c "
  nohup python -m vllm.entrypoints.openai.api_server \
    --model \"$MODEL_NAME\" \
    --served-model-name \"$SERVED_NAME\" \
    --trust-remote-code \
    --dtype bfloat16 \
    --tensor-parallel-size \"$TP_SIZE\" \
    --gpu-memory-utilization \"$GPU_UTIL\" \
    --max-num-seqs \"$MAX_NUM_SEQS\" \
    --max-model-len \"$MAX_MODEL_LEN\" \
    --port \"$API_PORT\" \
    >\"$LOG_DIR/vllm_${MODEL_SIZE}.log\" 2>&1 &
  echo \$! > \"$VLLM_PID_FILE\"
  sleep 0.5
" >/dev/null 2>&1 &

sleep 1
VLLM_PID="$(cat "$VLLM_PID_FILE" 2>/dev/null || true)"
if [[ -n "${VLLM_PID:-}" ]]; then
  VLLM_PGID="$(ps -o pgid= "$VLLM_PID" 2>/dev/null | tr -d ' ' || true)"
  [[ -n "${VLLM_PGID:-}" ]] && echo "$VLLM_PGID" > "$VLLM_PGID_FILE"
fi

cleanup() {
  set +e
  echo "[CLEANUP] Triggered. Cleaning vLLM..."
  if [[ -f "$VLLM_PGID_FILE" ]]; then
    pgid="$(cat "$VLLM_PGID_FILE" 2>/dev/null || true)"
    if [[ -n "${pgid:-}" ]]; then
      echo "[CLEANUP] kill -TERM -- -$pgid"
      kill -TERM -- "-$pgid" 2>/dev/null || true
      for _ in {1..15}; do lsof -iTCP:${API_PORT} -sTCP:LISTEN -Pn >/dev/null 2>&1 || break; sleep 1; done
      if lsof -iTCP:${API_PORT} -sTCP:LISTEN -Pn >/dev/null 2>&1; then
        echo "[CLEANUP] kill -KILL -- -$pgid"
        kill -KILL -- "-$pgid" 2>/dev/null || true
      fi
    fi
  fi
  if [[ -f "$VLLM_PID_FILE" ]]; then
    pid="$(cat "$VLLM_PID_FILE" 2>/dev/null || true)"
    if [[ -n "${pid:-}" ]] && ps -p "$pid" >/dev/null 2>&1; then
      kill -TERM "$pid" 2>/dev/null || true
      for _ in {1..5}; do ps -p "$pid" >/dev/null 2>&1 || break; sleep 1; done
      ps -p "$pid" >/dev/null 2>&1 && kill -KILL "$pid" 2>/dev/null || true
    fi
  fi
  if lsof -iTCP:${API_PORT} -sTCP:LISTEN -Pn >/dev/null 2>&1; then
    old_pids=$(lsof -iTCP:${API_PORT} -sTCP:LISTEN -t || true)
    [[ -n "$old_pids" ]] && kill $old_pids 2>/dev/null || true
  fi
  rm -f "$VLLM_PID_FILE" "$VLLM_PGID_FILE"
}
trap 'cleanup; exit 130' INT
trap 'cleanup' TERM EXIT

# 等待就绪
echo "[4/6] Waiting for vLLM to be ready at ${API_BASE} (timeout=${WAIT_SECS}s) ..."
ready=0
for (( i=1; i<=WAIT_SECS; i++ )); do
  if curl -m 2 -s "${API_BASE}/models" >/dev/null 2>&1; then ready=1; break; fi
  sleep 1
done
[[ $ready -eq 1 ]] || { echo "ERROR: vLLM not ready after ${WAIT_SECS}s. See $LOG_DIR/vllm_${MODEL_SIZE}.log"; exit 1; }
echo "vLLM is ready."

# 健康检查（标准 text）
if ! curl -m 5 -s -X POST "${API_BASE}/chat/completions" \
    -H 'Content-Type: application/json' \
    -d "{\"model\":\"${SERVED_NAME}\",\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"text\",\"text\":\"ping\"}]}],\"max_tokens\":4}" \
    | grep -q '"choices"'; then
  echo "WARNING: /v1/chat/completions 返回异常。确认 OpenAI 兼容入口与依赖是否齐全。"
fi

# 跑流水线（把 served-model-name 传入）
echo "[5/6] Running pipeline..."
CMD=(python "$ROOT_DIR/run.py"
  --csv "$CSV_PATH"
  --outdir "$OUTDIR"
  --api-base "$API_BASE"
  --model "$SERVED_NAME"
  --results-csv "$RESULTS_CSV"
)

# 只有等于 "1" 才加对应 flag；"0" 或空都不加
if [[ "${RESUME:-0}" == "1" ]]; then
  CMD+=("--resume")
fi
if [[ "${DRY_RUN:-0}" == "1" ]]; then
  CMD+=("--dry-run")
fi

echo "Command: ${CMD[*]}"
"${CMD[@]}"

echo "[6/6] Done."
echo "结果汇总 CSV: $RESULTS_CSV"
echo "日志目录    : $LOG_DIR"