#!/usr/bin/env bash
# verify-stepfun-cn.sh
# StepFun 中国区 Step Plan + Pi Provider 综合验证脚本
#
# 默认测试：
#   1. 可选：Pi 扩展加载与模型注册
#   2. step-3.7-flash 的 low / medium / high
#   3. step-3.5-flash-2603 的 low / high，并探测 medium
#   4. step-3.5-flash 无 reasoning_effort
#   5. step-router-v1 基础对话，并探测 low / medium / high
#   6. step-3.7-flash 与 step-router-v1 的两轮 function tool call
#   7. 可选：指定字符规模的上下文探测
#
# 注意：
#   - 该脚本验证“实际接口行为”，不能证明官方最大 contextWindow/maxTokens 极限。
#   - 上下文探测按字符数生成，不等于精确 Token 数。
#   - 不会打印 API Key。

set -Eeuo pipefail

BASE_URL="${STEPFUN_BASE_URL:-https://api.stepfun.com/step_plan/v1}"
EXTENSION_PATH=""
CONTEXT_CHARS=0
MAX_TOKENS=512

PASS_COUNT=0
WARN_COUNT=0
FAIL_COUNT=0
ROUTER_LOW=0
ROUTER_MEDIUM=0
ROUTER_HIGH=0

if [[ -t 1 ]]; then
  RED='\033[0;31m'
  GREEN='\033[0;32m'
  YELLOW='\033[1;33m'
  BLUE='\033[0;34m'
  NC='\033[0m'
else
  RED=''
  GREEN=''
  YELLOW=''
  BLUE=''
  NC=''
fi

pass() {
  PASS_COUNT=$((PASS_COUNT + 1))
  printf '%b[PASS]%b %s\n' "$GREEN" "$NC" "$*"
}

warn() {
  WARN_COUNT=$((WARN_COUNT + 1))
  printf '%b[WARN]%b %s\n' "$YELLOW" "$NC" "$*"
}

fail() {
  FAIL_COUNT=$((FAIL_COUNT + 1))
  printf '%b[FAIL]%b %s\n' "$RED" "$NC" "$*"
}

info() {
  printf '%b[INFO]%b %s\n' "$BLUE" "$NC" "$*"
}

usage() {
  cat <<'USAGE'
用法：
  export STEPFUN_API_KEY="你的中国区 Step Plan Key"
  bash verify-stepfun-cn.sh [选项]

选项：
  --extension PATH       临时加载该 stepfun-cn.ts，并检查 Pi 模型注册
  --context-chars N      可选上下文探测，生成 N 个字符（不是 Token）
  --max-tokens N         普通测试最大输出，默认 512
  --base-url URL         覆盖 API 地址
  -h, --help             显示帮助

示例：
  bash verify-stepfun-cn.sh \
    --extension ./stepfun-cn.ts

  bash verify-stepfun-cn.sh \
    --extension ./stepfun-cn.ts \
    --context-chars 131072
USAGE
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --extension)
      EXTENSION_PATH="${2:-}"
      shift 2
      ;;
    --context-chars)
      CONTEXT_CHARS="${2:-0}"
      shift 2
      ;;
    --max-tokens)
      MAX_TOKENS="${2:-512}"
      shift 2
      ;;
    --base-url)
      BASE_URL="${2:-}"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "未知参数：$1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

for cmd in curl jq python3; do
  if ! command -v "$cmd" >/dev/null 2>&1; then
    echo "缺少依赖：$cmd" >&2
    exit 2
  fi
done

if [[ -z "${STEPFUN_API_KEY:-}" ]]; then
  echo "未设置 STEPFUN_API_KEY。" >&2
  echo '请先运行：read -s "STEPFUN_API_KEY?StepFun API Key: "; echo; export STEPFUN_API_KEY' >&2
  exit 2
fi

if ! [[ "$CONTEXT_CHARS" =~ ^[0-9]+$ ]]; then
  echo "--context-chars 必须是非负整数" >&2
  exit 2
fi

if ! [[ "$MAX_TOKENS" =~ ^[1-9][0-9]*$ ]]; then
  echo "--max-tokens 必须是正整数" >&2
  exit 2
fi

TMP_DIR="$(mktemp -d)"
trap 'rm -rf "$TMP_DIR"' EXIT

info "API 地址：$BASE_URL"
info "STEPFUN_API_KEY 已设置（长度：${#STEPFUN_API_KEY}，不会打印内容）"

safe_error() {
  local file="$1"
  jq '{
    error: .error,
    model: .model,
    finish_reason: .choices[0].finish_reason,
    content: .choices[0].message.content
  }' "$file" 2>/dev/null || head -c 500 "$file"
  printf '\n'
}

post_payload() {
  local payload_file="$1"
  local response_file="$2"
  local http_code

  if ! http_code="$({
    curl -sS \
      -o "$response_file" \
      -w '%{http_code}' \
      -X POST "$BASE_URL/chat/completions" \
      -H 'Content-Type: application/json' \
      -H "Authorization: Bearer $STEPFUN_API_KEY" \
      --data-binary "@$payload_file"
  } 2>"$TMP_DIR/curl-error.log")"; then
    printf 'curl 请求失败：%s\n' "$(cat "$TMP_DIR/curl-error.log")" >&2
    return 1
  fi

  printf '%s' "$http_code"
}

build_text_payload() {
  local model="$1"
  local effort="$2"
  local expected="$3"
  local output_file="$4"

  if [[ -n "$effort" ]]; then
    jq -n \
      --arg model "$model" \
      --arg effort "$effort" \
      --arg expected "$expected" \
      --argjson max_tokens "$MAX_TOKENS" '
      {
        model: $model,
        messages: [
          {
            role: "user",
            content: ("请只回复 " + $expected)
          }
        ],
        reasoning_effort: $effort,
        max_tokens: $max_tokens
      }
    ' >"$output_file"
  else
    jq -n \
      --arg model "$model" \
      --arg expected "$expected" \
      --argjson max_tokens "$MAX_TOKENS" '
      {
        model: $model,
        messages: [
          {
            role: "user",
            content: ("请只回复 " + $expected)
          }
        ],
        max_tokens: $max_tokens
      }
    ' >"$output_file"
  fi
}

text_test() {
  local model="$1"
  local effort="$2"
  local expected="$3"
  local mode="$4"   # strict | probe
  local label="$model"
  local payload="$TMP_DIR/payload-${model//\//_}-${effort:-none}.json"
  local response="$TMP_DIR/response-${model//\//_}-${effort:-none}.json"
  local http_code content finish returned_model has_reasoning

  if [[ -n "$effort" ]]; then
    label="$label / reasoning_effort=$effort"
  else
    label="$label / 无 reasoning_effort"
  fi

  build_text_payload "$model" "$effort" "$expected" "$payload"

  if ! http_code="$(post_payload "$payload" "$response")"; then
    [[ "$mode" == "strict" ]] && fail "${label}：请求未完成" || warn "${label}：请求未完成"
    return 1
  fi

  if [[ "$http_code" != "200" ]]; then
    if [[ "$mode" == "strict" ]]; then
      fail "${label}：HTTP ${http_code}"
    else
      warn "${label}：HTTP ${http_code}（探测结果，不计为核心失败）"
    fi
    safe_error "$response"
    return 1
  fi

  content="$(jq -r '.choices[0].message.content // ""' "$response")"
  finish="$(jq -r '.choices[0].finish_reason // ""' "$response")"
  returned_model="$(jq -r '.model // ""' "$response")"
  has_reasoning="$(jq -r '((.choices[0].message.reasoning // "") != "" or (.choices[0].message.reasoning_content // "") != "")' "$response")"

  if [[ "$finish" == "length" ]]; then
    if [[ "$mode" == "strict" ]]; then
      fail "${label}：finish_reason=length，输出被截断"
    else
      warn "${label}：finish_reason=length，探测不完整"
    fi
    return 1
  fi

  if printf '%s' "$content" | grep -Fqi "$expected"; then
    pass "${label}：返回 ${expected}；returned_model=${returned_model:-未知}；reasoning=${has_reasoning}；finish=${finish}"
    return 0
  fi

  if [[ "$mode" == "strict" ]]; then
    fail "${label}：HTTP 200，但最终 content 不含 ${expected}"
  else
    warn "${label}：HTTP 200，但最终 content 不含 ${expected}"
  fi
  safe_error "$response"
  return 1
}

validate_extension() {
  local output="$TMP_DIR/pi-models.txt"
  # Pi 当前的 --list-models 按模型 ID 列出条目，不保证输出 provider/model。
  local expected_models=(
    'step-3.7-flash'
    'step-3.5-flash-2603'
    'step-3.5-flash'
    'step-router-v1'
  )
  local model

  [[ -z "$EXTENSION_PATH" ]] && return 0

  if ! command -v pi >/dev/null 2>&1; then
    fail "指定了 --extension，但系统找不到 pi 命令"
    return 1
  fi

  if [[ ! -f "$EXTENSION_PATH" ]]; then
    fail "扩展文件不存在：$EXTENSION_PATH"
    return 1
  fi

  if pi -e "$EXTENSION_PATH" --list-models >"$output" 2>&1; then
    pass "Pi 成功临时加载扩展：$EXTENSION_PATH"
  else
    fail "Pi 无法加载扩展：$EXTENSION_PATH"
    sed -n '1,120p' "$output"
    return 1
  fi

  for model in "${expected_models[@]}"; do
    if grep -Fq "$model" "$output"; then
      pass "Pi 模型已注册：${model}"
    else
      fail "Pi 模型未注册：${model}"
    fi
  done

  if grep -Eq 'oauth:|onPrompt|refreshToken|getApiKey' "$EXTENSION_PATH"; then
    pass "扩展源码包含 /login 所需 OAuth 认证结构"
  else
    warn "未在扩展源码中发现完整 OAuth 关键词；请人工确认 /login"
  fi
}

tool_loop_test() {
  local model="$1"
  local effort="$2"
  local label="$model 工具调用"
  local first_payload="$TMP_DIR/tool-first-${model}.json"
  local first_response="$TMP_DIR/tool-first-response-${model}.json"
  local second_payload="$TMP_DIR/tool-second-${model}.json"
  local second_response="$TMP_DIR/tool-second-response-${model}.json"
  local http_code tool_count tool_name tool_call_id assistant_message final_content finish

  if [[ -n "$effort" ]]; then
    label="$label / reasoning_effort=$effort"
  fi

  jq -n \
    --arg model "$model" \
    --arg effort "$effort" '
    {
      model: $model,
      messages: [
        {
          role: "user",
          content: "必须调用 list_project_files 工具，参数 path 必须为 .；不要直接编造文件列表。"
        }
      ],
      tools: [
        {
          type: "function",
          function: {
            name: "list_project_files",
            description: "列出项目目录中的文件",
            parameters: {
              type: "object",
              properties: {
                path: {type: "string"}
              },
              required: ["path"],
              additionalProperties: false
            }
          }
        }
      ],
      tool_choice: "auto",
      max_tokens: 1024
    }
    + if $effort != "" then {reasoning_effort: $effort} else {} end
  ' >"$first_payload"

  if ! http_code="$(post_payload "$first_payload" "$first_response")"; then
    fail "${label}：首轮请求失败"
    return 1
  fi

  if [[ "$http_code" != "200" ]]; then
    fail "${label}：首轮 HTTP ${http_code}"
    safe_error "$first_response"
    return 1
  fi

  tool_count="$(jq -r '(.choices[0].message.tool_calls // []) | length' "$first_response")"
  tool_name="$(jq -r '.choices[0].message.tool_calls[0].function.name // ""' "$first_response")"
  tool_call_id="$(jq -r '.choices[0].message.tool_calls[0].id // ""' "$first_response")"

  if [[ "$tool_count" -lt 1 || "$tool_name" != "list_project_files" || -z "$tool_call_id" ]]; then
    fail "${label}：未返回预期 function tool call"
    safe_error "$first_response"
    return 1
  fi

  pass "${label}：首轮成功生成 tool call（${tool_name}）"

  assistant_message="$(jq '.choices[0].message' "$first_response")"

  jq -n \
    --arg model "$model" \
    --arg effort "$effort" \
    --arg tool_call_id "$tool_call_id" \
    --argjson assistant "$assistant_message" '
    {
      model: $model,
      messages: [
        {
          role: "user",
          content: "必须调用 list_project_files 工具，参数 path 必须为 .；不要直接编造文件列表。"
        },
        $assistant,
        {
          role: "tool",
          tool_call_id: $tool_call_id,
          content: "README.md\nstepfun-cn.ts\ntest.sh"
        }
      ],
      tools: [
        {
          type: "function",
          function: {
            name: "list_project_files",
            description: "列出项目目录中的文件",
            parameters: {
              type: "object",
              properties: {
                path: {type: "string"}
              },
              required: ["path"],
              additionalProperties: false
            }
          }
        }
      ],
      tool_choice: "auto",
      max_tokens: 1024
    }
    + if $effort != "" then {reasoning_effort: $effort} else {} end
  ' >"$second_payload"

  if ! http_code="$(post_payload "$second_payload" "$second_response")"; then
    fail "${label}：第二轮请求失败"
    return 1
  fi

  if [[ "$http_code" != "200" ]]; then
    fail "${label}：第二轮 HTTP ${http_code}"
    safe_error "$second_response"
    return 1
  fi

  final_content="$(jq -r '.choices[0].message.content // ""' "$second_response")"
  finish="$(jq -r '.choices[0].finish_reason // ""' "$second_response")"

  if printf '%s' "$final_content" | grep -Fqi 'README.md'; then
    pass "${label}：两轮 Agent tool loop 成功；finish=${finish}"
    return 0
  fi

  fail "${label}：第二轮未基于工具结果给出最终回答"
  safe_error "$second_response"
  return 1
}

context_probe() {
  local chars="$1"
  local marker="CTX_OK_${chars}"
  local content_file="$TMP_DIR/context.txt"
  local payload="$TMP_DIR/context-payload.json"
  local response="$TMP_DIR/context-response.json"
  local http_code content finish

  [[ "$chars" -le 0 ]] && return 0

  info "开始上下文探测：${chars} 个字符；这不是精确 Token 测试"

  python3 - "$chars" "$marker" >"$content_file" <<'PY'
import sys
count = int(sys.argv[1])
marker = sys.argv[2]
prefix = "A" * max(0, count - len(marker) - 20)
print(prefix)
print("末尾标记：" + marker)
PY

  jq -n \
    --arg model 'step-3.7-flash' \
    --arg marker "$marker" \
    --rawfile context "$content_file" '
    {
      model: $model,
      messages: [
        {
          role: "user",
          content: ($context + "\n请只回复末尾标记：" + $marker)
        }
      ],
      reasoning_effort: "low",
      max_tokens: 256
    }
  ' >"$payload"

  if ! http_code="$(post_payload "$payload" "$response")"; then
    fail "上下文探测：请求失败"
    return 1
  fi

  if [[ "$http_code" != "200" ]]; then
    fail "上下文探测：HTTP $http_code"
    safe_error "$response"
    return 1
  fi

  content="$(jq -r '.choices[0].message.content // ""' "$response")"
  finish="$(jq -r '.choices[0].finish_reason // ""' "$response")"

  if printf '%s' "$content" | grep -Fq "$marker"; then
    pass "上下文探测通过：${chars} 字符；finish=$finish"
  else
    fail "上下文探测未返回末尾标记：${chars} 字符"
    safe_error "$response"
  fi
}

printf '\n===== 1. Pi 扩展注册检查 =====\n'
validate_extension || true

printf '\n===== 2. 固定模型文本与思考档位 =====\n'
text_test 'step-3.7-flash' 'low' 'STEP37_LOW_OK' strict || true
text_test 'step-3.7-flash' 'medium' 'STEP37_MEDIUM_OK' strict || true
text_test 'step-3.7-flash' 'high' 'STEP37_HIGH_OK' strict || true

text_test 'step-3.5-flash-2603' 'low' 'STEP35_2603_LOW_OK' strict || true
text_test 'step-3.5-flash-2603' 'high' 'STEP35_2603_HIGH_OK' strict || true
text_test 'step-3.5-flash-2603' 'medium' 'STEP35_2603_MEDIUM_PROBE' probe || true

text_test 'step-3.5-flash' '' 'STEP35_OK' strict || true

printf '\n===== 3. Step Router V1 =====\n'
text_test 'step-router-v1' '' 'ROUTER_BASE_OK' strict || true
if text_test 'step-router-v1' 'low' 'ROUTER_LOW_OK' probe; then ROUTER_LOW=1; fi
if text_test 'step-router-v1' 'medium' 'ROUTER_MEDIUM_OK' probe; then ROUTER_MEDIUM=1; fi
if text_test 'step-router-v1' 'high' 'ROUTER_HIGH_OK' probe; then ROUTER_HIGH=1; fi

printf '\n===== 4. Function tool call 两轮循环 =====\n'
tool_loop_test 'step-3.7-flash' 'low' || true
tool_loop_test 'step-router-v1' '' || true

printf '\n===== 5. 可选上下文探测 =====\n'
context_probe "$CONTEXT_CHARS" || true

printf '\n===== Router 配置建议 =====\n'
if [[ "$ROUTER_LOW" -eq 1 && "$ROUTER_MEDIUM" -eq 1 && "$ROUTER_HIGH" -eq 1 ]]; then
  pass 'Router 的 low / medium / high 均实测成功：可考虑 supportsReasoningEffort=true，并开放三档。'
elif [[ "$ROUTER_LOW" -eq 0 && "$ROUTER_MEDIUM" -eq 0 && "$ROUTER_HIGH" -eq 0 ]]; then
  warn 'Router 三个 reasoning_effort 探测均未通过：建议 supportsReasoningEffort=false，不配置 thinkingLevelMap。'
else
  warn "Router reasoning_effort 支持不一致：low=${ROUTER_LOW} medium=${ROUTER_MEDIUM} high=${ROUTER_HIGH}；只开放稳定通过的档位，或保守设为 false。"
fi

printf '\n===== 最终汇总 =====\n'
printf 'PASS: %d\nWARN: %d\nFAIL: %d\n' "$PASS_COUNT" "$WARN_COUNT" "$FAIL_COUNT"

if [[ "$FAIL_COUNT" -gt 0 ]]; then
  printf '%b结论：存在核心失败，请先修复后再正式安装 Provider。%b\n' "$RED" "$NC"
  exit 1
fi

printf '%b结论：核心接口测试通过；WARN 项属于能力探测或文档边界，需要按结果决定 Provider 元数据。%b\n' "$GREEN" "$NC"
exit 0