#!/usr/bin/env python3
"""One-off diagnostic: see exactly what z.ai's Anthropic-compatible endpoint
returns, so the ZaiProvider's rate-limit header mapping can be finalized against
a real response (the header names aren't documented).

Usage (key via env so it never lands in shell history):
    ZAI_API_KEY=... python tools/probe_zai.py [base_url] [model]

Defaults: base = https://api.z.ai/api/anthropic , model = glm-4.6

Prints the HTTP status and the FULL response header set (highlighting any header
mentioning limit/rate/quota/usage/reset), plus a short body preview. Sends a
1-token "hi" — on a prompt-metered plan that is one prompt, so run it sparingly.
The API key is read from the environment and never printed.
"""
import json
import os
import sys

import httpx

BASE = (sys.argv[1] if len(sys.argv) > 1 else "https://api.z.ai/api/anthropic").rstrip("/")
MODEL = sys.argv[2] if len(sys.argv) > 2 else "glm-4.6"
KEY = os.environ.get("ZAI_API_KEY") or os.environ.get("ANTHROPIC_AUTH_TOKEN")

if not KEY:
    sys.exit("Set ZAI_API_KEY in the environment first (it is not printed).")

url = f"{BASE}/v1/messages"
headers = {
    "Authorization": f"Bearer {KEY}",
    "anthropic-version": "2023-06-01",
    "Content-Type": "application/json",
}
body = {"model": MODEL, "max_tokens": 1, "messages": [{"role": "user", "content": "hi"}]}

print(f"POST {url}  (model={MODEL})")
try:
    resp = httpx.post(url, headers=headers, json=body, timeout=30.0)
except httpx.HTTPError as e:
    sys.exit(f"request failed: {e}")

print(f"HTTP {resp.status_code}\n")
INTEREST = ("limit", "rate", "quota", "usage", "reset", "remaining")
print("--- response headers ---")
for k in sorted(resp.headers.keys()):
    mark = "  <<<" if any(w in k.lower() for w in INTEREST) else ""
    print(f"  {k}: {resp.headers[k]}{mark}")

print("\n--- body preview ---")
try:
    print(json.dumps(resp.json(), ensure_ascii=False)[:600])
except ValueError:
    print(resp.text[:600])