"""unifyllm 网关连通性 / tool-calling 冒烟测试。

用法(.env 在仓库根, 含 UNIFYLLM_API_KEY):
    .venv/Scripts/python.exe scripts/diag_unifyllm.py [model_id ...]

不带参数时测默认 5 个精选模型。走 litellm openai/ 前缀 + api_base,
与 core/llm.py 线上路径一致。输出一律 ASCII(Windows 控制台 GBK)。
"""
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
env_file = os.path.join(ROOT, ".env")
if os.path.exists(env_file):
    for line in open(env_file, encoding="utf-8"):
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            k, v = line.split("=", 1)
            os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))

os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
import litellm  # noqa: E402

API_BASE = "https://unifyllm.ai/v1"
DEFAULT_MODELS = [
    "claude-fable-5",
    "claude-opus-4-8",
    "claude-sonnet-4-6",
    "gpt-5.6-sol",
    "gemini-3.1-pro-preview",
]

TOOLS = [{
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get current weather for a city",
        "parameters": {
            "type": "object",
            "properties": {"city": {"type": "string"}},
            "required": ["city"],
        },
    },
}]


def main() -> int:
    key = os.environ.get("UNIFYLLM_API_KEY")
    if not key:
        print("[FAIL] UNIFYLLM_API_KEY not set (check .env)")
        return 1
    args = sys.argv[1:]
    force_temp = None
    stream = False
    for a in list(args):
        if a.startswith("--temp="):
            force_temp = float(a.split("=", 1)[1])
            args.remove(a)
        elif a == "--stream":
            stream = True
            args.remove(a)
    models = args or DEFAULT_MODELS
    failed = []
    for m in models:
        # litellm 客户端硬拦 gpt-5* 的 temperature != 1
        temp = force_temp if force_temp is not None else (1.0 if m.startswith("gpt-5") else 0.3)
        try:
            kwargs = dict(
                model=f"openai/{m}",
                api_base=API_BASE,
                api_key=key,
                messages=[{"role": "user", "content": "What's the weather in Beijing? Use the tool."}],
                tools=TOOLS,
                temperature=temp,
                timeout=120,
            )
            if stream:
                # 与 core/llm.py chat_stream 同路径:流式 + include_usage,chunk 拼回完整 response
                chunks = list(litellm.completion(**kwargs, stream=True,
                                                 stream_options={"include_usage": True}))
                resp = litellm.stream_chunk_builder(chunks)
            else:
                resp = litellm.completion(**kwargs)
            msg = resp.choices[0].message
            tc = msg.tool_calls or []
            usage = resp.usage
            if tc:
                print(f"[OK]   {m}: tool_call={tc[0].function.name}({tc[0].function.arguments.strip()}) "
                      f"tokens={usage.prompt_tokens}+{usage.completion_tokens}")
            else:
                text = (msg.content or "")[:80].replace("\n", " ")
                print(f"[WARN] {m}: no tool_call, text={text!r}")
        except Exception as e:
            failed.append(m)
            print(f"[FAIL] {m}: {type(e).__name__}: {str(e)[:200]}")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())