"""Deepseek API 流式快速验证脚本。

按照需求：
- deepseek-chat: max_tokens=8000
- deepseek-reasoner: max_tokens=65536
- 开启流式输出，打印每条 data 行便于确认返回。

根据用户允许，直接在脚本中硬编码 DeepSeek 测试用 API 信息，避免被仓库中其他
默认配置（如 kimi）覆盖。请勿在生产使用。
"""

from __future__ import annotations

import asyncio
from typing import Optional

import httpx


DEEPSEEK_BASE_URL = "https://api.deepseek.com"
DEEPSEEK_API_KEY = "sk-3457fbc33f0b4aefb2ce1d3101bb2341"
DEEPSEEK_CHAT_MODEL = "deepseek-chat"
DEEPSEEK_REASONER_MODEL = "deepseek-reasoner"


def build_headers(api_key: str) -> dict[str, str]:
    return {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }


async def stream_call(
    *,
    name: str,
    base_url: str,
    api_key: str,
    model: str,
    max_tokens: int,
    prompt: Optional[str] = None,
) -> None:
    """向指定模型发起流式 chat/completions 请求并打印 data 行。"""
    url = base_url.rstrip("/") + "/chat/completions"
    payload = {
        "model": model,
        "stream": True,
        "max_tokens": max_tokens,
        "messages": [
            {
                "role": "user",
                "content": prompt
                or f"这是 {name} 模型的流式测试，请用一句话自我介绍。",
            }
        ],
    }

    print(f"\n=== {name} ===")
    print(f"POST {url}")
    async with httpx.AsyncClient(http2=True, timeout=120) as client:
        async with client.stream(
            "POST",
            url,
            json=payload,
            headers=build_headers(api_key),
        ) as resp:
            print(f"status: {resp.status_code}")
            if resp.status_code != 200:
                body = await resp.aread()
                print("error body:", body.decode(errors="ignore"))
                return
            async for line in resp.aiter_lines():
                if not line:
                    continue
                if line.startswith("data:"):
                    data = line[5:].strip()
                    if data == "[DONE]":
                        print("[DONE]")
                        break
                    print(data)
                else:
                    # 兼容潜在的非 data 行（例如心跳）
                    print(line)


async def main() -> None:
    await stream_call(
        name="deepseek-chat (max_tokens=8000)",
        base_url=DEEPSEEK_BASE_URL,
        api_key=DEEPSEEK_API_KEY,
        model=DEEPSEEK_CHAT_MODEL,
        max_tokens=8000,
    )

    await stream_call(
        name="deepseek-reasoner (max_tokens=65536)",
        base_url=DEEPSEEK_BASE_URL,
        api_key=DEEPSEEK_API_KEY,
        model=DEEPSEEK_REASONER_MODEL,
        max_tokens=65536,
        prompt="你是一个思考模型，请简述测试状态并结束。",
    )


if __name__ == "__main__":
    asyncio.run(main())