#!/usr/bin/env bash
#
# Llama 4 统一容器入口：
#   1. 幂等应用 CoreX vLLM 的 Llama 4 补丁
#   2. 读取 /models/llama4/config.json，自动判断 Scout / Maverick
#   3. 选择对应的 TP / PP / 端口 / served name，再启动 vLLM
#
# AUTO_DETECT=1（默认）时根据模型 config 自动选择；显式设置的
# TP / PP / PORT / SERVED_NAME 会覆盖检测结果。AUTO_DETECT=0 可完全手动。
#
set -euo pipefail

PATCH_SCRIPT="${LLAMA4_PATCH_SCRIPT:-/opt/llama4/apply_patch.py}"
DETECT_SCRIPT="${LLAMA4_DETECT_SCRIPT:-/opt/llama4/detect_model.py}"
VLLM_BIN="${VLLM_BIN:-vllm}"

python3 "${PATCH_SCRIPT}"

MODEL_PATH="${MODEL_PATH:-/models/llama4}"
AUTO_DETECT="${AUTO_DETECT:-1}"
MODEL_DESCRIPTION=""

if [[ ! -f "${MODEL_PATH}/config.json" ]]; then
  echo "ERROR: 模型目录不完整：${MODEL_PATH}/config.json 不存在" >&2
  echo "请确认已挂载模型，例如：" >&2
  echo "  -v /mnt/disk0/models/LLM-Research/Llama-4-Scout-17B-16E-Instruct:/models/llama4:ro" >&2
  echo "  -v /mnt/disk0/models/LLM-Research/Llama-4-Maverick-17B-128E-Instruct-FP8:/models/llama4:ro" >&2
  exit 1
fi

if [[ "${AUTO_DETECT}" == "1" ]]; then
  DETECT_OUTPUT="$(python3 "${DETECT_SCRIPT}" "${MODEL_PATH}")" || exit 1
  eval "${DETECT_OUTPUT}"
  MODEL_KIND="${MODEL_KIND:-${DETECTED_MODEL_KIND}}"
  TP="${TP:-${DETECTED_TP}}"
  PP="${PP:-${DETECTED_PP}}"
  MODEL_DESCRIPTION="${DETECTED_DESCRIPTION:-}"
  # 检测出的 port / served_name 仅用于日志提示，不再决定容器内监听端口和模型名：
  # 容器内端口固定 80，模型名固定 llm（见 AGENTS.md 5.2 / 5.5）。
  DETECTED_HOST_PORT_HINT="${DETECTED_PORT:-}"
  DETECTED_NAME_HINT="${DETECTED_SERVED_NAME:-}"
else
  MODEL_KIND="${MODEL_KIND:-unknown}"
  DETECTED_HOST_PORT_HINT=""
  DETECTED_NAME_HINT=""
fi

HOST="${HOST:-0.0.0.0}"
TP="${TP:-8}"
PP="${PP:-1}"
PORT="${PORT:-80}"
SERVED_NAME="${SERVED_NAME:-llm}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-1}"
GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.90}"
ENFORCE_EAGER="${ENFORCE_EAGER:-1}"
EXTRA_ARGS="${EXTRA_ARGS:-}"

echo "============================================================"
echo "Llama 4 (BI-V150 unified image)"
echo "  model kind : ${MODEL_KIND}"
echo "  model      : ${MODEL_PATH}"
echo "  served name: ${SERVED_NAME}"
echo "  listen     : ${HOST}:${PORT}"
echo "  TP / PP    : ${TP} / ${PP}  (total $((TP * PP)) GPUs)"
echo "  max len    : ${MAX_MODEL_LEN}"
echo "  max seqs   : ${MAX_NUM_SEQS}"
echo "  gpu util   : ${GPU_MEM_UTIL}"
echo "  eager      : ${ENFORCE_EAGER}"
echo "  visible    : ${CUDA_VISIBLE_DEVICES:-all}"
echo "  extra      : ${EXTRA_ARGS}"
if [[ -n "${MODEL_DESCRIPTION}" ]]; then
  echo "  detected   : ${MODEL_DESCRIPTION}"
fi
if [[ -n "${DETECTED_HOST_PORT_HINT}" ]]; then
  echo "  hint       : 变体默认宿主端口 ${DETECTED_HOST_PORT_HINT}，模型名曾用 ${DETECTED_NAME_HINT}"
  echo "               （容器内固定 80 / served name 固定 llm，实际端口由 docker run -p 决定）"
fi
echo "============================================================"

# ---------------------------------------------------------------- 用户参数
# 优先级（高 -> 低）：
#   1. docker run <image> 之后追加的 vllm 参数（"$@"）
#   2. -e EXTRA_ARGS="..."
#   3. -e 显式设置（SERVED_NAME / PORT / TP / PP / MAX_MODEL_LEN ...）
#   4. config.json 自动检测
# 用户已经给了的 flag，本脚本不再重复传，避免 vllm 报 duplicate key。
# 末尾参数的 flag 优先于 EXTRA_ARGS 里的同名 flag：把后者摘掉，避免 vllm 报 duplicate key。
declare -A _CLI_FLAGS=()
if (( $# > 0 )); then
  for _t in "$@"; do
    [[ "${_t}" == --* ]] && _CLI_FLAGS["${_t%%=*}"]=1
  done
fi

EXTRA_FILTERED=()
if [[ -n "${EXTRA_ARGS}" ]]; then
  read -r -a _extra_tokens <<< "${EXTRA_ARGS}"
  _skip=0
  for _t in "${_extra_tokens[@]}"; do
    if (( _skip > 0 )); then
      _skip=0
      continue
    fi
    if [[ "${_t}" == --* ]]; then
      _name="${_t%%=*}"
      if [[ -n "${_CLI_FLAGS[${_name}]:-}" ]]; then
        # 该 flag 在命令行里也给了，跳过它（以及紧跟的取值）
        [[ "${_t}" != *=* ]] && _skip=1
        continue
      fi
    fi
    EXTRA_FILTERED+=("${_t}")
  done
fi

USER_TOKENS=("${EXTRA_FILTERED[@]}")
if (( $# > 0 )); then
  USER_TOKENS+=("$@")
fi

user_has() {
  local want="$1" tok
  (( ${#USER_TOKENS[@]} == 0 )) && return 1
  for tok in "${USER_TOKENS[@]}"; do
    if [[ "${tok}" == "${want}" || "${tok}" == "${want}="* ]]; then
      return 0
    fi
  done
  return 1
}

if (( $# > 0 )); then
  echo "用户追加参数 : $*"
  echo "（优先级最高，覆盖 -e 与自动检测）"
  user_has --port && echo "  WARN: --port 会改变容器内监听端口，宿主机 -p 的右侧端口也要同步改" >&2
  user_has --served-model-name && echo "  WARN: --served-model-name 已覆盖默认的 llm，客户端 model 字段需同步修改" >&2
  echo "------------------------------------------------------------"
fi

SERVE_ARGS=()
user_has --host || SERVE_ARGS+=(--host "${HOST}")
user_has --port || SERVE_ARGS+=(--port "${PORT}")
user_has --served-model-name || SERVE_ARGS+=(--served-model-name "${SERVED_NAME}")
user_has --tensor-parallel-size || SERVE_ARGS+=(--tensor-parallel-size "${TP}")
if (( PP > 1 )) && ! user_has --pipeline-parallel-size; then
  SERVE_ARGS+=(--pipeline-parallel-size "${PP}")
fi
user_has --max-model-len || SERVE_ARGS+=(--max-model-len "${MAX_MODEL_LEN}")
user_has --max-num-seqs || SERVE_ARGS+=(--max-num-seqs "${MAX_NUM_SEQS}")
user_has --gpu-memory-utilization || SERVE_ARGS+=(--gpu-memory-utilization "${GPU_MEM_UTIL}")
if ! user_has --enforce-eager && ! user_has --no-enforce-eager; then
  if [[ "${ENFORCE_EAGER}" == "1" ]]; then
    SERVE_ARGS+=(--enforce-eager)
  fi
fi

# shellcheck disable=SC2086
exec "${VLLM_BIN}" serve "${MODEL_PATH}" \
  "${SERVE_ARGS[@]}" \
  "${EXTRA_FILTERED[@]}" \
  "$@"