#!/bin/bash

read -p "请输入 Base URL: " base_url
read -p "请输入 API Key: " api_key

base_url=${base_url%/}
models_url="${base_url}/models"
chat_url="${base_url}/chat/completions"

models_id=()
models_name=()

while true; do
    echo -e "\n=== API 检测与测试菜单 ==="
    echo "1. 获取模型列表"
    echo "2. 使用编号进行聊天测试 (带自动重试)"
    echo "3. 退出"
    read -p "请选择 (1-3): " choice

    case $choice in
        1)
            models_id=(); models_name=()
            while IFS= read -r line; do
                models_id+=("$(echo "$line" | cut -d'|' -f1)")
                models_name+=("$(echo "$line" | cut -d'|' -f2)")
            done < <(curl -s -H "Authorization: Bearer $api_key" "$models_url" | jq -r '.data[] | "\(.id)|\(.name // .id)"')
            for i in "${!models_id[@]}"; do printf "%d) ID: %-30s | Name: %s\n" "$((i+1))" "${models_id[$i]}" "${models_name[$i]}"; done
            ;;
        2)
            if [ ${#models_id[@]} -eq 0 ]; then echo "请先执行选项 1！"; continue; fi
            read -p "请输入模型编号: " idx
            if [[ "$idx" =~ ^[0-9]+$ ]] && [ "$idx" -ge 1 ] && [ "$idx" -le "${#models_id[@]}" ]; then
                selected_model=${models_id[$((idx-1))]}
                read -p "请输入消息: " user_content
                
                max_retries=2
                attempt=0
                success=false
                
                while [ $attempt -le $max_retries ]; do
                    echo "正在尝试 (第 $((attempt+1)) 次)..."
                    response=$(curl -s "$chat_url" \
                        -H "Content-Type: application/json" \
                        -H "Authorization: Bearer $api_key" \
                        -d "{
                            \"model\": \"$selected_model\",
                            \"messages\": [{\"role\": \"user\", \"content\": \"$user_content\"}],
                            \"max_tokens\": 100
                        }")
                    
                    reply=$(echo "$response" | jq -r '.choices[0].message.content // empty')
                    
                    if [ -n "$reply" ]; then
                        echo -e "\n模型回复:\n$reply"
                        success=true
                        break
                    else
                        echo "警告: 未收到有效回复。"
                        attempt=$((attempt+1))
                    fi
                done
                
                [ "$success" = false ] && echo "错误: 经过多次尝试，依然无法获取回复。"
            else
                echo "无效编号。"
            fi
            ;;
        3) break ;;
        *) echo "无效选择。" ;;
    esac
done