#!/bin/bash
# 豆包 API 连通性测试脚本
# 用法: bash test-doubao-api.sh

API_KEY="18480c26-ebcd-4263-8a6f-48359b8bd65d"
BASE_URL="https://ark.cn-beijing.volces.com/api/v3"
MODEL="doubao-seed-2-0-pro-260215"

echo "============================================"
echo "  豆包（火山引擎 Ark）API 测试"
echo "============================================"
echo ""

# ---- 测试 1: 非流式调用 ----
echo "【测试 1】非流式调用"
echo "请求: 你好，请用一句话介绍你自己"
echo "---"

START=$(date +%s%3N 2>/dev/null || python3 -c "import time; print(int(time.time()*1000))")

RESPONSE=$(curl -s -w "\n---HTTP_CODE:%{http_code}---" \
  -X POST "${BASE_URL}/chat/completions" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer ${API_KEY}" \
  -d "{
    \"model\": \"${MODEL}\",
    \"messages\": [{\"role\": \"user\", \"content\": \"你好，请用一句话介绍你自己\"}],
    \"stream\": false
  }" 2>&1)

END=$(date +%s%3N 2>/dev/null || python3 -c "import time; print(int(time.time()*1000))")
ELAPSED=$((END - START))

HTTP_CODE=$(echo "$RESPONSE" | grep -oP '(?<=---HTTP_CODE:)\d+(?=---)')
BODY=$(echo "$RESPONSE" | sed 's/---HTTP_CODE:[0-9]*---//')

echo "HTTP 状态码: ${HTTP_CODE}"
echo "耗时: ${ELAPSED} ms"

if [ "$HTTP_CODE" = "200" ]; then
    # 提取回复内容
    CONTENT=$(echo "$BODY" | python3 -c "
import sys, json
data = json.load(sys.stdin)
choices = data.get('choices', [])
if choices:
    print('回复:', choices[0].get('message', {}).get('content', '(empty)'))
usage = data.get('usage', {})
print(f'Token: prompt={usage.get(\"prompt_tokens\",0)}, completion={usage.get(\"completion_tokens\",0)}, total={usage.get(\"total_tokens\",0)}')
print(f'模型: {data.get(\"model\", \"?\")}')
" 2>&1)
    echo "$CONTENT"
else
    echo "错误响应: $BODY"
fi

echo ""
echo ""

# ---- 测试 2: 流式调用 ----
echo "【测试 2】流式调用 (SSE)"
echo "请求: 简单说一下健康饮食的三个要点"
echo "---"

START2=$(date +%s%3N 2>/dev/null || python3 -c "import time; print(int(time.time()*1000))")

# 流式请求，逐行输出
STREAM_OUTPUT=""
FIRST_CHUNK_TIME=""

curl -s -N \
  -X POST "${BASE_URL}/chat/completions" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer ${API_KEY}" \
  -d "{
    \"model\": \"${MODEL}\",
    \"messages\": [{\"role\": \"user\", \"content\": \"简单说一下健康饮食的三个要点，每个要点一句话\"}],
    \"stream\": true
  }" 2>/dev/null | while IFS= read -r line; do
    if [[ "$line" == data:* ]]; then
        JSON_DATA="${line#data: }"
        if [ "$JSON_DATA" = "[DONE]" ]; then
            break
        fi
        # 提取 delta content
        DELTA=$(echo "$JSON_DATA" | python3 -c "
import sys, json
try:
    d = json.load(sys.stdin)
    c = d.get('choices',[{}])[0].get('delta',{}).get('content','')
    if c: print(c, end='', flush=True)
except: pass
" 2>/dev/null)
        if [ -z "$FIRST_CHUNK_TIME" ] && [ -n "$DELTA" ]; then
            FIRST_CHUNK_TIME=$(date +%s%3N 2>/dev/null || python3 -c "import time; print(int(time.time()*1000))")
        fi
    fi
done

END2=$(date +%s%3N 2>/dev/null || python3 -c "import time; print(int(time.time()*1000))")
ELAPSED2=$((END2 - START2))

echo ""
echo "---"
echo "流式总耗时: ${ELAPSED2} ms"

echo ""
echo ""

# ---- 测试 3: 通过后端接口调用（需先启动后端服务） ----
echo "【测试 3】通过后端接口调用（需后端服务运行中）"
echo "请确认后端已启动，默认地址: http://localhost:8081"
echo "---"

BACKEND_URL="http://localhost:8081/api/front/doubao/chat"

RESP3=$(curl -s -w "\n---HTTP_CODE:%{http_code}---" \
  -X POST "${BACKEND_URL}" \
  -H "Content-Type: application/json" \
  -d "{
    \"messages\": [{\"role\": \"user\", \"content\": \"你好\"}]
  }" 2>&1)

HTTP3=$(echo "$RESP3" | grep -oP '(?<=---HTTP_CODE:)\d+(?=---)')
BODY3=$(echo "$RESP3" | sed 's/---HTTP_CODE:[0-9]*---//')

if [ "$HTTP3" = "200" ]; then
    echo "后端接口调用成功! HTTP $HTTP3"
    echo "$BODY3" | python3 -c "
import sys, json
data = json.load(sys.stdin)
choices = data.get('choices', [])
if choices:
    print('回复:', choices[0].get('message', {}).get('content', '(empty)')[:100])
" 2>&1
else
    echo "后端接口调用失败: HTTP ${HTTP3:-连接失败}"
    echo "（如果后端未启动，这是正常的）"
fi

echo ""
echo "============================================"
echo "  测试完成"
echo "============================================"