#!/usr/bin/env python3
"""count_tokens_example.py — 用官方 Token Counting API 精确计数（需要 API key）

【需 key】本脚本调用真实 API（POST /v1/messages/count_tokens）。
没有 ANTHROPIC_API_KEY 时会打印提示并退出；脚本已做静态校验：
- 请求结构符合官方 API 参考
- key 从环境变量读取，绝不硬编码

来源:
- https://platform.claude.com/docs/en/api/messages-count-tokens（访问于 2026-08-30）
- https://platform.claude.com/docs/en/build-with-claude/token-counting（访问于 2026-08-30）

先设置环境变量 ANTHROPIC_API_KEY，再运行: python3 count_tokens_example.py
"""

import json
import os
import sys
import urllib.request

API_URL = "https://api.anthropic.com/v1/messages/count_tokens"
API_VERSION = "2023-06-01"


def build_request() -> dict:
    """构造 count_tokens 请求体：与 Messages API 相同的 model/system/messages。"""
    return {
        "model": "claude-sonnet-5",
        "system": "你是企业知识库助手，回答必须引用来源文档段落。",
        "messages": [
            {"role": "user", "content": "总结 2026 年 Q2 销售数据的主要发现。"},
            {"role": "assistant", "content": "Q2 收入同比增长 12%，主要由欧洲区贡献。"},
            {"role": "user", "content": "欧洲区的具体数字是多少？"},
        ],
    }


def main() -> int:
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print("[需 key] 未检测到 ANTHROPIC_API_KEY。")
        print("设置后先设置环境变量 ANTHROPIC_API_KEY，再运行: python3 count_tokens_example.py")
        print("\n[静态校验] 请求体结构预览（可直接用于 curl 对照官方参考）:")
        print(json.dumps(build_request(), ensure_ascii=False, indent=2))
        return 0

    body = json.dumps(build_request()).encode()
    req = urllib.request.Request(
        API_URL,
        data=body,
        headers={
            "x-api-key": api_key,
            "anthropic-version": API_VERSION,
            "content-type": "application/json",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        result = json.loads(resp.read())
    print(f"精确 input_tokens = {result.get('input_tokens')}")
    print("把这个值乘以你的对话轮数增长速度，即可推导上下文耗尽时间点。")
    return 0


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