#!/usr/bin/env bash
# 读取 task_info.env，修补模型目录，并准备 vLLM 启动命令。
# --dry-run 只打印计划，不启动进程，不访问平台。
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

usage() {
  echo "用法: entrypoint.sh [--dry-run]" >&2
}

DRY_RUN=0
if [[ "${1:-}" == "--dry-run" ]]; then
  DRY_RUN=1
elif [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
  usage
  exit 0
elif [[ -n "${1:-}" ]]; then
  usage
  exit 1
fi

info_file="${TASK_INFO_FILE:-}"
if [[ -z "$info_file" ]]; then
  if [[ -f "${ROOT}/task_info.env" ]]; then
    info_file="${ROOT}/task_info.env"
  else
    info_file="${ROOT}/task_info.env.example"
  fi
fi

if [[ ! -f "$info_file" ]]; then
  echo "找不到任务信息文件: ${info_file}" >&2
  exit 1
fi

# 已在环境中的变量优先于文件，方便本地覆盖。
while IFS= read -r line || [[ -n "$line" ]]; do
  case "$line" in
    ""|\#*) continue ;;
  esac
  key="${line%%=*}"
  value="${line#*=}"
  if [[ -z "$key" || "$key" == "$line" ]]; then
    echo "任务信息行无法解析: ${line}" >&2
    exit 1
  fi
  if [[ -z "${!key+x}" ]]; then
    export "${key}=${value}"
  fi
done <"$info_file"

FRAMEWORK="${FRAMEWORK:-}"
GPU_TYPE="${GPU_TYPE:-}"
TASK_TYPE="${TASK_TYPE:-text-generation}"
MODEL_DIR="${MODEL_DIR:-/model}"
SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-llm}"
PORT="${PORT:-8000}"
ARCH_UNKNOWN_POLICY="${ARCH_UNKNOWN_POLICY:-fail}"
APPLY_TOKENIZER_PATCH="${APPLY_TOKENIZER_PATCH:-0}"
FORCE_CHAT_TEMPLATE="${FORCE_CHAT_TEMPLATE:-0}"

if [[ -z "$FRAMEWORK" || "$FRAMEWORK" == "vllm_fix_tokenizer" ]]; then
  echo "拒绝 FRAMEWORK=${FRAMEWORK:-空}。本仓使用 vllm_bi150_r2r4。" >&2
  exit 1
fi
if [[ -z "$GPU_TYPE" ]]; then
  echo "GPU_TYPE 为空" >&2
  exit 1
fi

plan=(
  python3 -m vllm.entrypoints.openai.api_server
  --model "$MODEL_DIR"
  --port "$PORT"
  --served-model-name "$SERVED_MODEL_NAME"
  --trust-remote-code
)

echo "FRAMEWORK=${FRAMEWORK}"
echo "GPU_TYPE=${GPU_TYPE}"
echo "TASK_TYPE=${TASK_TYPE}"
echo "MODEL_DIR=${MODEL_DIR}"
echo "ARCH_UNKNOWN_POLICY=${ARCH_UNKNOWN_POLICY}"
echo "plan: ${plan[*]}"

if [[ "$DRY_RUN" -eq 1 ]]; then
  preview="$(
    python3 - "$ROOT" <<'PY'
import sys
from pathlib import Path
root = Path(sys.argv[1])
sys.path.insert(0, str(root / "scripts"))
import apply_patches
messages = [
    {"role": "system", "content": "你是助手"},
    {"role": "user", "content": "你好"},
]
print(apply_patches.preview_default_template(messages).replace("\n", "\\n"))
PY
  )"
  echo "template_preview: ${preview}"
  echo "dry-run: 不启动进程，不访问平台"
  exit 0
fi

if [[ ! -d "$MODEL_DIR" ]]; then
  echo "模型目录不存在: ${MODEL_DIR}" >&2
  exit 1
fi

patch_args=(
  --model-dir "$MODEL_DIR"
  --framework "$FRAMEWORK"
  --unknown-policy "$ARCH_UNKNOWN_POLICY"
)
if [[ "$FORCE_CHAT_TEMPLATE" == "1" ]]; then
  patch_args+=(--force-chat-template)
fi
if [[ "$APPLY_TOKENIZER_PATCH" == "1" ]]; then
  patch_args+=(--tokenizer-patch)
fi

python3 "${ROOT}/scripts/apply_patches.py" "${patch_args[@]}"
exec "${plan[@]}"