"""
Sora2API 测试脚本
用于诊断长提示词问题
"""

import asyncio
import httpx
import json
from pathlib import Path
import tomllib

async def test_sora_api():
    """测试 Sora API"""

    # 读取配置
    config_path = Path(__file__).parent / "config.toml"
    with open(config_path, "rb") as f:
        config = tomllib.load(f)

    api_config = config["api"]

    # 测试用例
    test_cases = [
        {
            "name": "短提示词",
            "prompt": "一个美丽的日落场景",
            "model": "sora-video-portrait-15s"
        },
        {
            "name": "中等提示词",
            "prompt": "一个美丽的日落场景，天空中有橙色和粉色的云彩，海浪轻轻拍打着沙滩，远处有几只海鸥在飞翔",
            "model": "sora-video-portrait-15s"
        },
        {
            "name": "长提示词",
            "prompt": "之前去饭店，点了鳕鱼，商家拿油鱼冒充，吃完感觉有点腻，上车把老婆刚买的西梅汁一口干了，当时觉得挺解腻的。结果开车没多久，肚子就开始不对劲，咕噜咕噜响，然后就是一阵阵的绞痛。我赶紧找了个加油站，冲进厕所，那场面简直了，像开闸放水一样，根本控制不住。更尴尬的是，裤子都来不及脱，直接就喷了，油鱼的油脂混着西梅汁，那个味道，简直让人窒息。",
            "model": "sora-video-portrait-15s"
        }
    ]

    for test_case in test_cases:
        print(f"\n{'='*60}")
        print(f"测试: {test_case['name']}")
        print(f"提示词长度: {len(test_case['prompt'])} 字符")
        print(f"{'='*60}")

        try:
            url = f"{api_config['base_url'].rstrip('/')}/v1/chat/completions"
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_config['token']}"
            }

            payload = {
                "model": test_case["model"],
                "messages": [{"role": "user", "content": test_case["prompt"]}],
                "stream": True
            }

            print(f"请求URL: {url}")
            print(f"Payload大小: {len(json.dumps(payload, ensure_ascii=False))} 字节")
            print(f"开始请求...")

            timeout = httpx.Timeout(connect=10.0, read=60.0, write=10.0, pool=10.0)

            async with httpx.AsyncClient(timeout=timeout) as client:
                async with client.stream("POST", url, json=payload, headers=headers) as response:
                    print(f"响应状态码: {response.status_code}")
                    print(f"响应头: {dict(response.headers)}")

                    if response.status_code == 200:
                        print("开始接收流式响应...")
                        line_count = 0
                        async for line in response.aiter_lines():
                            line_count += 1
                            if line.startswith("data: "):
                                data_str = line[6:]
                                if data_str == "[DONE]":
                                    print(f"收到 [DONE] 标记")
                                    break
                                try:
                                    data = json.loads(data_str)
                                    if "choices" in data and data["choices"]:
                                        delta = data["choices"][0].get("delta", {})
                                        content = delta.get("content", "")
                                        if content:
                                            print(f"收到内容片段: {content[:100]}...")
                                except json.JSONDecodeError as e:
                                    print(f"JSON解析失败: {e}, 数据: {data_str[:100]}...")

                        print(f"✅ 测试成功！共收到 {line_count} 行响应")
                    else:
                        error_text = await response.aread()
                        print(f"❌ 请求失败: {response.status_code}")
                        print(f"错误信息: {error_text.decode('utf-8', errors='ignore')[:500]}")

        except Exception as e:
            import traceback
            print(f"❌ 异常: {type(e).__name__}: {str(e)}")
            print(f"详细信息:\n{traceback.format_exc()}")

        # 等待一下再进行下一个测试
        await asyncio.sleep(2)

    print(f"\n{'='*60}")
    print("所有测试完成")
    print(f"{'='*60}")

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