"""
Qwen API 连通性测试脚本
验证 API Key 和模型名称是否可用
"""
import os
import sys

# 手动读取 .env（避免引入项目依赖）
env_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
if os.path.exists(env_file):
    for line in open(env_file):
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            key, _, val = line.partition("=")
            os.environ.setdefault(key.strip(), val.strip())

API_KEY = os.environ.get("QWEN_API_KEY", "")
BASE_URL = os.environ.get("QWEN_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
MODEL = os.environ.get("QWEN_MODEL", "qwen-plus")
MAX_TOKENS = int(os.environ.get("QWEN_MAX_TOKENS", "4096"))
TEMPERATURE = float(os.environ.get("QWEN_TEMPERATURE", "0.1"))

BLUE = "\033[34m"
GREEN = "\033[32m"
RED = "\033[31m"
YELLOW = "\033[33m"
RESET = "\033[0m"

print(f"\n{BLUE}╔══════════════════════════════════════╗{RESET}")
print(f"{BLUE}║     Qwen API 连通性测试               ║{RESET}")
print(f"{BLUE}╚══════════════════════════════════════╝{RESET}\n")

# 检查配置
print(f"{YELLOW}[配置]{RESET}")
print(f"  API Key:  {API_KEY[:16]}...{API_KEY[-4:]}" if len(API_KEY) > 20 else f"  API Key:  {API_KEY or '(空)'}")
print(f"  Base URL: {BASE_URL}")
print(f"  Model:    {MODEL}")
print(f"  MaxTokens:{MAX_TOKENS}")
print(f"  Temp:     {TEMPERATURE}\n")

if not API_KEY:
    print(f"{RED}[错误] QWEN_API_KEY 为空，请先填入 .env 文件{RESET}")
    sys.exit(1)

# 引入 openai
try:
    from openai import OpenAI
except ImportError:
    print(f"{RED}[错误] 缺少 openai 包，请先安装: pip install openai{RESET}")
    sys.exit(1)

print(f"{YELLOW}[测试1] 基础连通性 — 发送简单问候...{RESET}")

client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

messages = [
    {"role": "system", "content": "你是一个药典助手，用中文简短回答。"},
    {"role": "user", "content": "你好，请用一句话介绍自己。"},
]

try:
    resp = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        max_tokens=256,
        temperature=TEMPERATURE,
    )
    answer = resp.choices[0].message.content or ""
    usage = resp.usage

    print(f"{GREEN}  ✓ 调用成功！{RESET}")
    print(f"  回答: {answer[:200]}")
    print(f"  Token用量: 输入={usage.prompt_tokens}, 输出={usage.completion_tokens}")
    print(f"  Model: {resp.model}\n")

except Exception as e:
    print(f"{RED}  ✗ 失败: {e}{RESET}\n")
    sys.exit(1)

# 测试2：模拟药典查询场景（测试长回答能力）
print(f"{YELLOW}[测试2] 药典场景 — 查询药品信息...{RESET}")

messages2 = [
    {
        "role": "system",
        "content": (
            "你是中华药典AI助手，知识来源为《中华人民共和国药典》。"
            "回答须基于药典原文，不得编造。回答格式：结论 → 原文引用 → 来源标注。"
        ),
    },
    {
        "role": "user",
        "content": "请模拟回答：阿莫西林的贮藏条件是什么？（即使没有检索到原文，也要模拟格式示范）",
    },
]

try:
    resp2 = client.chat.completions.create(
        model=MODEL,
        messages=messages2,
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE,
    )
    answer2 = resp2.choices[0].message.content or ""
    usage2 = resp2.usage

    print(f"{GREEN}  ✓ 调用成功！{RESET}")
    print(f"  回答: {answer2[:300]}...")
    print(f"  Token用量: 输入={usage2.prompt_tokens}, 输出={usage2.completion_tokens}")
    print(f"  Model: {resp2.model}\n")

except Exception as e:
    print(f"{RED}  ✗ 失败: {e}{RESET}\n")

# 总结
print(f"{BLUE}══════════════════════════════════════{RESET}")
print(f"{GREEN}  API Key 验证通过，Qwen API 可用！{RESET}")
print(f"{GREEN}  下一步：docker compose up -d 启动全栈环境{RESET}")
print(f"{BLUE}══════════════════════════════════════{RESET}\n")