#!/usr/bin/env python3
"""
豆包（火山引擎 Ark）API 快速测试
用法: python3 test-doubao-api.py
"""
import json, time, sys

try:
    import requests
except ImportError:
    print("需要安装 requests: pip install requests")
    sys.exit(1)

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

headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {API_KEY}"
}

print("=" * 50)
print("  豆包 API 测试")
print("=" * 50)

# ---- 测试 1: 非流式 ----
print("\n【测试 1】非流式调用")
payload = {
    "model": MODEL,
    "messages": [{"role": "user", "content": "你好，请用一句话介绍你自己"}],
    "stream": False
}

start = time.time()
try:
    resp = requests.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload, timeout=30)
    elapsed = (time.time() - start) * 1000
    print(f"HTTP 状态码: {resp.status_code}")
    print(f"耗时: {elapsed:.0f} ms")
    if resp.status_code == 200:
        data = resp.json()
        choices = data.get("choices", [])
        if choices:
            print(f"回复: {choices[0]['message']['content']}")
        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', '?')}")
    else:
        print(f"错误: {resp.text[:500]}")
except Exception as e:
    elapsed = (time.time() - start) * 1000
    print(f"请求失败 ({elapsed:.0f}ms): {e}")

# ---- 测试 2: 流式 ----
print("\n【测试 2】流式调用 (SSE)")
payload_stream = {
    "model": MODEL,
    "messages": [{"role": "user", "content": "简单说一下健康饮食的三个要点，每个要点一句话"}],
    "stream": True
}

start2 = time.time()
first_token_time = None
full_content = ""

try:
    resp2 = requests.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload_stream, stream=True, timeout=30)
    print(f"HTTP 状态码: {resp2.status_code}")
    if resp2.status_code == 200:
        print("流式输出: ", end="", flush=True)
        for line in resp2.iter_lines(decode_unicode=True):
            if not line:
                continue
            if line.startswith("data: "):
                json_str = line[6:].strip()
                if json_str == "[DONE]":
                    break
                try:
                    evt = json.loads(json_str)
                    delta = evt.get("choices", [{}])[0].get("delta", {}).get("content", "")
                    if delta:
                        if first_token_time is None:
                            first_token_time = time.time()
                        full_content += delta
                        print(delta, end="", flush=True)
                except json.JSONDecodeError:
                    pass

        total_time = (time.time() - start2) * 1000
        ttft = (first_token_time - start2) * 1000 if first_token_time else 0
        print(f"\n---")
        print(f"首字时间 (TTFT): {ttft:.0f} ms")
        print(f"总耗时: {total_time:.0f} ms")
        print(f"回复长度: {len(full_content)} 字符")
    else:
        print(f"错误: {resp2.text[:500]}")
except Exception as e:
    elapsed2 = (time.time() - start2) * 1000
    print(f"\n请求失败 ({elapsed2:.0f}ms): {e}")

print("\n" + "=" * 50)
print("  测试完成")
print("=" * 50)