#!/usr/bin/env bash
# test.sh - StepFun 中国区 Step Plan API 连通性测试
#
# 用途：
#   1. 检查 STEPFUN_API_KEY 环境变量是否存在（不输出 Key 内容）
#   2. 向 step-3.7-flash 发送请求，要求模型回复 "OK"
#   3. 返回明确退出码（0=成功，1=失败）
#
# 用法：
#   export STEPFUN_API_KEY="your-key"
#   bash test.sh
#
# 退出码：
#   0 - 测试通过（API 可达，content 包含大小写不敏感的 OK）
#   1 - 测试失败（环境变量未设置、API 不可达、或 content 不包含 OK）

set -euo pipefail

PROVIDER_BASE_URL="https://api.stepfun.com/step_plan/v1"
MODEL="step-3.7-flash"

# ---------- 颜色输出（可选，不影响退出码）----------
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

pass() { echo -e "${GREEN}[PASS]${NC} $1"; }
fail() { echo -e "${RED}[FAIL]${NC} $1"; }
info() { echo -e "${YELLOW}[INFO]${NC} $1"; }

# ---------- 1. 检查环境变量 ----------
if [ -z "${STEPFUN_API_KEY:-}" ]; then
  fail "环境变量 STEPFUN_API_KEY 未设置"
  info "请先设置: export STEPFUN_API_KEY=\"your-api-key\""
  exit 1
fi

info "STEPFUN_API_KEY 已设置 (长度: ${#STEPFUN_API_KEY} 字符)"
pass "环境变量检查通过"

# ---------- 2. 发送测试请求 ----------
info "正在请求 ${PROVIDER_BASE_URL}/chat/completions (模型: ${MODEL}) ..."

HTTP_RESPONSE=$(curl -sS -w "\n%{http_code}" \
  -X POST "${PROVIDER_BASE_URL}/chat/completions" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer ${STEPFUN_API_KEY}" \
  -d "{
    \"model\": \"${MODEL}\",
    \"messages\": [{\"role\": \"user\", \"content\": \"请只回复 OK 两个字\"}],
    \"max_tokens\": 256,
    \"reasoning_effort\": \"low\"
  }")

HTTP_CODE=$(echo "$HTTP_RESPONSE" | tail -n1)
BODY=$(echo "$HTTP_RESPONSE" | sed '$d')

if [ "$HTTP_CODE" != "200" ]; then
  fail "HTTP 状态码: ${HTTP_CODE}"
  info "响应体（截断）: ${BODY:0:200}"
  exit 1
fi

pass "HTTP 请求成功 (状态码: ${HTTP_CODE})"

# ---------- 3. 解析响应 ----------
# 读取 .choices[0].message.content 和 .choices[0].finish_reason
# reasoning / reasoning_content 仅用于故障诊断，不作为通过依据
REPLY_TEXT=$(echo "$BODY" | python3 -c "
import sys, json
d = json.load(sys.stdin)
choice = d.get('choices', [{}])[0]
content = choice.get('message', {}).get('content', '')
finish_reason = choice.get('finish_reason', '')
print(content if content else '')
print(finish_reason)
" 2>/dev/null)

REPLY_CONTENT=$(echo "$REPLY_TEXT" | head -n1)
FINISH_REASON=$(echo "$REPLY_TEXT" | tail -n1)

# 故障诊断：输出 reasoning / reasoning_content（不用于判断通过/失败）
REASONING_DIAG=$(echo "$BODY" | python3 -c "
import sys, json
d = json.load(sys.stdin)
msg = d.get('choices', [{}])[0].get('message', {})
rc = msg.get('reasoning_content', '')
reasoning = msg.get('reasoning', '')
if rc:
    print('reasoning_content: ' + rc[:120])
if reasoning:
    print('reasoning: ' + reasoning[:120])
" 2>/dev/null)

if [ -n "$REASONING_DIAG" ]; then
  info "诊断信息: ${REASONING_DIAG}"
fi

info "finish_reason: ${FINISH_REASON}"
info "content: ${REPLY_CONTENT}"

# ---------- 4. 判断结果 ----------
# 只有 content 包含大小写不敏感的 OK 才算通过
# reasoning / reasoning_content 仅用于故障诊断，不作为通过依据

if [ -z "$REPLY_CONTENT" ] && [ "$FINISH_REASON" = "length" ]; then
  fail "content 为空且 finish_reason=length，token 不足"
  info "请增加 max_tokens 值后重试"
  exit 1
fi

if echo "$REPLY_CONTENT" | grep -qi "OK"; then
  pass "content 包含 'OK'，测试通过！"
  exit 0
else
  fail "content 不包含 'OK'（内容: ${REPLY_CONTENT}）"
  info "reasoning / reasoning_content 仅供参考，不作为通过依据"
  exit 1
fi