#!/bin/bash

# SGLang 多模型 curl 测试脚本
# 测试模型：Qwen3-8B, Qwen3.5-35B, Qwen3.5-122B, Qwen3-Embedding-8B, Qwen3-Reranker-8B

set -e

# 颜色定义
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 模型配置
declare -A MODELS=(
    ["qwen3-8b"]="25424|/model/Qwen3-8B|chat"
    ["qwen3.5-35b"]="25427|/model/Qwen3.5-35B|chat"
    ["qwen3.5-122b"]="25423|/model/Qwen3.5-122B-A10B|chat"
    ["qwen3-embedding-8b"]="25425|/model/Qwen3-Embedding-8B|embedding"
    ["qwen3-reranker-8b"]="25426|/model/Qwen3-Reranker-8B|rerank"
)

API_KEY="lq123456"
TIMEOUT=30

echo "========================================"
echo "SGLang 多模型健康检查 (curl)"
echo "时间: $(date '+%Y-%m-%d %H:%M:%S')"
echo "========================================"

# 测试计数器
TOTAL=0
SUCCESS=0

# 测试对话模型
test_chat_model() {
    local name=$1
    local port=$2
    local model_path=$3
    
    echo ""
    echo "----------------------------------------"
    echo "测试模型: $name (对话模型)"
    echo "端口: $port"
    echo "----------------------------------------"
    
    local response
    response=$(curl -s -w "\n%{http_code}" \
        --max-time $TIMEOUT \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "{
            \"model\": \"$model_path\",
            \"messages\": [{\"role\": \"user\", \"content\": \"你好，请用一句话介绍自己\"}],
            \"temperature\": 0.7,
            \"max_tokens\": 50
        }" \
        "http://localhost:$port/v1/chat/completions" 2>/dev/null || echo -e "\n000")
    
    local body=$(echo "$response" | head -n -1)
    local code=$(echo "$response" | tail -n 1)
    
    if [ "$code" = "200" ]; then
        local content=$(echo "$body" | grep -o '"content":"[^"]*"' | head -1 | cut -d'"' -f4)
        echo -e "${GREEN}✅ 成功${NC} HTTP $code"
        echo "回复: ${content:0:100}..."
        ((SUCCESS++))
    else
        echo -e "${RED}❌ 失败${NC} HTTP $code"
        echo "响应: ${body:0:200}"
    fi
    ((TOTAL++))
}

# 测试嵌入模型
test_embedding_model() {
    local name=$1
    local port=$2
    local model_path=$3
    
    echo ""
    echo "----------------------------------------"
    echo "测试模型: $name (嵌入模型)"
    echo "端口: $port"
    echo "----------------------------------------"
    
    local response
    response=$(curl -s -w "\n%{http_code}" \
        --max-time $TIMEOUT \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "{
            \"model\": \"$model_path\",
            \"input\": [\"你好，这是一个测试句子\", \"Hello world\"]
        }" \
        "http://localhost:$port/v1/embeddings" 2>/dev/null || echo -e "\n000")
    
    local body=$(echo "$response" | head -n -1)
    local code=$(echo "$response" | tail -n 1)
    
    if [ "$code" = "200" ]; then
        local dims=$(echo "$body" | grep -o '"embedding":\[[^]]*\]' | head -1 | grep -o ',' | wc -l)
        dims=$((dims + 1))
        echo -e "${GREEN}✅ 成功${NC} HTTP $code"
        echo "向量维度: $dims"
        ((SUCCESS++))
    else
        echo -e "${YELLOW}⚠️  Embedding 接口失败，尝试 Rerank 接口...${NC}"
        # 尝试 rerank 接口
        response=$(curl -s -w "\n%{http_code}" \
            --max-time $TIMEOUT \
            -H "Content-Type: application/json" \
            -H "Authorization: Bearer $API_KEY" \
            -d "{
                \"model\": \"$model_path\",
                \"query\": \"测试查询\",
                \"documents\": [\"文档1\", \"文档2\"]
            }" \
            "http://localhost:$port/v1/rerank" 2>/dev/null || echo -e "\n000")
        
        code=$(echo "$response" | tail -n 1)
        if [ "$code" = "200" ]; then
            echo -e "${GREEN}✅ 成功${NC} (Rerank 接口可用)"
            ((SUCCESS++))
        else
            echo -e "${RED}❌ 失败${NC} HTTP $code"
        fi
    fi
    ((TOTAL++))
}

# 测试重排序模型
test_rerank_model() {
    local name=$1
    local port=$2
    local model_path=$3
    
    echo ""
    echo "----------------------------------------"
    echo "测试模型: $name (重排序模型)"
    echo "端口: $port"
    echo "----------------------------------------"
    
    # 尝试 rerank 接口
    local response
    response=$(curl -s -w "\n%{http_code}" \
        --max-time $TIMEOUT \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $API_KEY" \
        -d "{
            \"model\": \"$model_path\",
            \"query\": \"什么是机器学习\",
            \"documents\": [\"机器学习是AI的分支\", \"Python是编程语言\", \"深度学习使用神经网络\"],
            \"top_n\": 2
        }" \
        "http://localhost:$port/v1/rerank" 2>/dev/null || echo -e "\n000")
    
    local body=$(echo "$response" | head -n -1)
    local code=$(echo "$response" | tail -n 1)
    
    if [ "$code" = "200" ]; then
        local top_doc=$(echo "$body" | grep -o '"text":"[^"]*"' | head -1 | cut -d'"' -f4)
        local score=$(echo "$body" | grep -o '"relevance_score":[0-9.]*' | head -1 | cut -d':' -f2)
        echo -e "${GREEN}✅ 成功${NC} HTTP $code"
        echo "Top1: ${top_doc:0:50}... (得分: $score)"
        ((SUCCESS++))
    else
        echo -e "${YELLOW}⚠️  Rerank 接口失败，尝试 Chat 接口...${NC}"
        # 尝试作为 chat 模型
        response=$(curl -s -w "\n%{http_code}" \
            --max-time $TIMEOUT \
            -H "Content-Type: application/json" \
            -H "Authorization: Bearer $API_KEY" \
            -d "{
                \"model\": \"$model_path\",
                \"messages\": [{\"role\": \"user\", \"content\": \"你好\"}],
                \"max_tokens\": 20
            }" \
            "http://localhost:$port/v1/chat/completions" 2>/dev/null || echo -e "\n000")
        
        code=$(echo "$response" | tail -n 1)
        if [ "$code" = "200" ]; then
            echo -e "${GREEN}✅ 成功${NC} (Chat 接口可用)"
            ((SUCCESS++))
        else
            echo -e "${RED}❌ 失败${NC} HTTP $code"
        fi
    fi
    ((TOTAL++))
}

# 快速检查服务是否存活
quick_check() {
    echo ""
    echo "========================================"
    echo "快速检查模式"
    echo "========================================"
    
    for key in "${!MODELS[@]}"; do
        IFS='|' read -r port model_path mtype <<< "${MODELS[$key]}"
        
        local code
        code=$(curl -s -o /dev/null -w "%{http_code}" \
            --max-time 5 \
            -H "Authorization: Bearer $API_KEY" \
            "http://localhost:$port/v1/models" 2>/dev/null || echo "000")
        
        if [ "$code" = "200" ]; then
            echo -e "${GREEN}✅${NC} $key (端口 $port)"
            ((SUCCESS++))
        else
            echo -e "${RED}❌${NC} $key (端口 $port) HTTP $code"
        fi
        ((TOTAL++))
    done
}

# 主函数
main() {
    # 解析参数
    if [ "$1" = "--quick" ]; then
        quick_check
    elif [ "$1" = "--model" ] && [ -n "$2" ]; then
        # 测试指定模型
        if [ -n "${MODELS[$2]}" ]; then
            IFS='|' read -r port model_path mtype <<< "${MODELS[$2]}"
            case $mtype in
                chat) test_chat_model "$2" "$port" "$model_path" ;;
                embedding) test_embedding_model "$2" "$port" "$model_path" ;;
                rerank) test_rerank_model "$2" "$port" "$model_path" ;;
            esac
        else
            echo "未知模型: $2"
            echo "可用模型: ${!MODELS[@]}"
            exit 1
        fi
    else
        # 完整测试所有模型
        for key in "${!MODELS[@]}"; do
            IFS='|' read -r port model_path mtype <<< "${MODELS[$key]}"
            case $mtype in
                chat) test_chat_model "$key" "$port" "$model_path" ;;
                embedding) test_embedding_model "$key" "$port" "$model_path" ;;
                rerank) test_rerank_model "$key" "$port" "$model_path" ;;
            esac
        done
    fi
    
    # 输出摘要
    echo ""
    echo "========================================"
    echo "测试结果摘要"
    echo "========================================"
    echo "总计: $SUCCESS / $TOTAL 个模型正常"
    
    if [ $SUCCESS -eq $TOTAL ]; then
        echo -e "${GREEN}所有模型运行正常！${NC}"
        exit 0
    else
        echo -e "${RED}部分模型异常，请检查日志${NC}"
        exit 1
    fi
}

# 执行
main "$@"