#!/bin/bash
set -e

MODEL_DIR=${1:-/model}
shift || true

FIX_TOKENIZER_DIR=/tmp/fixed_tokenizer
AUTO_FIX=${AUTO_FIX_TOKENIZER:-auto}

echo "[entrypoint] model dir: $MODEL_DIR"

NEED_FIX=0

if [ "$AUTO_FIX" = "1" ] || [ "$AUTO_FIX" = "true" ]; then
    NEED_FIX=1
elif [ "$AUTO_FIX" = "auto" ]; then
    if [ -f "$MODEL_DIR/tokenizer_config.json" ]; then
        if grep -q "TokenizersBackend\|TiktokenTokenizer" "$MODEL_DIR/tokenizer_config.json"; then
            NEED_FIX=1
        fi
        # 检测 extra_special_tokens 是否为 list 格式
        if grep -q '"extra_special_tokens":\s*\[' "$MODEL_DIR/tokenizer_config.json"; then
            NEED_FIX=1
        fi
    fi
fi

if [ $NEED_FIX -eq 1 ]; then
    echo "[entrypoint] fixing tokenizer..."
    python3 /opt/fix_tokenizer.py
    TOKENIZER_ARG="--tokenizer $FIX_TOKENIZER_DIR"
else
    echo "[entrypoint] tokenizer OK, skip fix"
    TOKENIZER_ARG=""
fi

# 检查并patch head_size支持（即使失败也不影响启动）
echo "[entrypoint] checking model head_size..."
python3 /opt/detect_head_size.py || echo "[entrypoint] head_size check failed, but continuing with vllm startup"

echo "[entrypoint] starting vllm..."

exec vllm serve "$MODEL_DIR" $TOKENIZER_ARG "$@"