#!/bin/bash
#
# DGX Spark vLLM モデル起動スクリプト
# Usage: ./dgx-spark-serve.sh [model] [options...]
#
set -euo pipefail

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info()  { echo -e "${BLUE}[INFO]${NC} $1"; }
log_ok()    { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# デフォルト設定
DEFAULT_TP_SIZE=2
DEFAULT_MAX_MODEL_LEN=65536
DEFAULT_GPU_MEM_UTIL=0.90
DEFAULT_HOST="0.0.0.0"
DEFAULT_PORT=8000

# プリセットモデル
declare -A MODEL_PRESETS=(
    ["minimax-m25"]="cerebras/MiniMax-M2.5-REAP-172B-A10B --tool-call-parser minimax_m2 --reasoning-parser minimax_m2_append_think --enable-auto-tool-choice"
    ["minimax-m2"]="cerebras/MiniMax-M2-REAP-172B-A10B --tool-call-parser minimax_m2 --reasoning-parser minimax_m2_append_think --enable-auto-tool-choice"
    ["llama-70b"]="nvidia/Llama-3.3-70B-Instruct-NVFP4"
    ["qwen-32b"]="nvidia/Qwen3-32B-NVFP4"
    ["gpt-oss-120b"]="openai/gpt-oss-120b"
)

show_presets() {
    echo "利用可能なプリセット:"
    for key in "${!MODEL_PRESETS[@]}"; do
        echo "  $key"
    done
}

find_container() {
    docker ps --format '{{.Names}}' | grep -E '^node-[0-9]+$' | head -1
}

main() {
    local model="${1:-}"
    shift || true
    
    if [[ -z "$model" || "$model" == "help" || "$model" == "--help" ]]; then
        cat << 'EOF'
DGX Spark vLLM モデル起動スクリプト

Usage:
  dgx-spark-serve.sh <model|preset> [options...]

Presets:
  minimax-m25   MiniMax-M2.5-REAP-172B (推奨)
  minimax-m2    MiniMax-M2-REAP-172B
  llama-70b     Llama-3.3-70B-Instruct-NVFP4
  qwen-32b      Qwen3-32B-NVFP4
  gpt-oss-120b  GPT-OSS-120B

Options (vllm serveに渡される):
  --tensor-parallel-size N  テンソル並列数 (default: 2)
  --max-model-len N         最大コンテキスト長 (default: 65536)
  --host IP                 APIホスト (default: 0.0.0.0)
  --port N                  APIポート (default: 8000)

Examples:
  # プリセット使用
  ./dgx-spark-serve.sh minimax-m25

  # カスタムモデル
  ./dgx-spark-serve.sh my-org/my-model --max-model-len 32768

  # メモリ節約モード
  ./dgx-spark-serve.sh minimax-m25 --max-model-len 16384 --max-num-seqs 32
EOF
        echo ""
        show_presets
        exit 0
    fi
    
    # コンテナ確認
    local container
    container=$(find_container)
    
    if [[ -z "$container" ]]; then
        log_error "vLLMコンテナが見つかりません"
        log_info "先に 'dgx-spark-setup.sh cluster' でクラスターを起動してください"
        exit 1
    fi
    
    log_ok "コンテナ検出: $container"
    
    # プリセット展開
    local model_args=""
    if [[ -n "${MODEL_PRESETS[$model]:-}" ]]; then
        model_args="${MODEL_PRESETS[$model]}"
        log_info "プリセット使用: $model"
        model=$(echo "$model_args" | awk '{print $1}')
        model_args=$(echo "$model_args" | cut -d' ' -f2-)
    fi
    
    # デフォルトオプション構築
    local has_tp=false has_len=false has_host=false has_port=false has_trust=false
    for arg in "$@"; do
        case "$arg" in
            --tensor-parallel-size*) has_tp=true ;;
            --max-model-len*) has_len=true ;;
            --host*) has_host=true ;;
            --port*) has_port=true ;;
            --trust-remote-code*) has_trust=true ;;
        esac
    done
    
    local defaults=""
    $has_tp    || defaults+=" --tensor-parallel-size $DEFAULT_TP_SIZE"
    $has_len   || defaults+=" --max-model-len $DEFAULT_MAX_MODEL_LEN"
    $has_host  || defaults+=" --host $DEFAULT_HOST"
    $has_port  || defaults+=" --port $DEFAULT_PORT"
    $has_trust || defaults+=" --trust-remote-code"
    defaults+=" --gpu-memory-utilization $DEFAULT_GPU_MEM_UTIL"
    
    # コマンド構築
    local cmd="vllm serve $model $model_args $defaults $*"
    
    log_info "起動コマンド:"
    echo "  $cmd"
    echo ""
    
    read -rp "実行しますか？ [Y/n]: " confirm
    if [[ "${confirm,,}" == "n" ]]; then
        log_warn "キャンセルしました"
        exit 0
    fi
    
    log_info "モデルを起動中..."
    docker exec -it "$container" /bin/bash -c "$cmd"
}

main "$@"