#!/usr/bin/env python3
"""token_estimation.py — 上下文预算估算器（无需 API key，本地可跑）

演示架构师做容量规划时的核心动作：
1. 用启发式估算文本 token 数（中/英/代码差异）
2. 计算一次请求的上下文预算占用（system + 历史 + 工具定义 + 输出预留）
3. 演示 context window 超限判断

注意：启发式估算只用于规划。精确值必须用官方 Token Counting API
（POST /v1/messages/count_tokens，见 count_tokens_example.py）。

来源: https://platform.claude.com/docs/en/build-with-claude/context-windows（访问于 2026-08-30）
运行: python3 token_estimation.py
"""

import math

# 启发式系数（粗估，用于规划；精确值以官方 count_tokens API 为准）：
# 英文约 4 字符/token（官方参考 1M token ≈ 4.5M 英文字符 ≈ 555k 单词），
# 中文常见汉字约 1.6 字符/token
EN_CHARS_PER_TOKEN = 4
ZH_CHARS_PER_TOKEN = 2.6

# 官方模型参数（截至 2026-08-30；易变事实，使用前请重验）
# 来源: https://platform.claude.com/docs/en/models/overview（访问于 2026-08-30）
MODELS = {
    "claude-opus-5":    {"context": 1_000_000, "max_output": 128_000},
    "claude-sonnet-5":  {"context": 1_000_000, "max_output": 128_000},
    "claude-haiku-4-5": {"context": 200_000, "max_output": 64_000},
}


def estimate_tokens(text: str) -> int:
    """极简启发式：按中英字符分别估算。生产环境请用 count_tokens API。"""
    zh = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
    other = len(text) - zh
    return math.ceil(zh / ZH_CHARS_PER_TOKEN + other / EN_CHARS_PER_TOKEN)


def image_tokens(width: int, height: int) -> int:
    """图片视觉 token：⌈宽/28⌉ × ⌈高/28⌉。
    来源: https://platform.claude.com/docs/en/build-with-claude/vision（访问于 2026-08-30）"""
    return math.ceil(width / 28) * math.ceil(height / 28)


def plan_context_budget(model: str, system: str, history_msgs: list[str],
                        tool_defs_tokens: int, output_reserve: int) -> None:
    spec = MODELS[model]
    window = spec["context"]

    system_t = estimate_tokens(system)
    history_t = sum(estimate_tokens(m) for m in history_msgs)
    image_t = image_tokens(1000, 1000)  # 一张 1000x1000 图
    used = system_t + history_t + tool_defs_tokens + image_t
    remaining = window - used - output_reserve

    print(f"=== 上下文预算规划: {model}（窗口 {window:,}，输出上限 {spec['max_output']:,}）===")
    print(f"  system prompt        : {system_t:>8,} tokens")
    print(f"  对话历史 ({len(history_msgs)} 条) : {history_t:>8,} tokens")
    print(f"  工具定义             : {tool_defs_tokens:>8,} tokens")
    print(f"  图片 1000x1000       : {image_t:>8,} tokens（⌈1000/28⌉²）")
    print(f"  输出预留             : {output_reserve:>8,} tokens")
    print(f"  {'-'*44}")
    print(f"  剩余可用             : {remaining:>8,} tokens")
    if remaining < 0:
        print("  [!] 超出窗口！应用层动作：压缩历史 / 检索裁剪 / 换更大窗口模型")
    elif remaining < window * 0.1:
        print("  [!] 剩余不足 10%：准备 compaction 或裁剪策略")
    else:
        print("  [OK] 预算充足")


if __name__ == "__main__":
    system = "你是企业知识库助手，回答必须引用来源文档段落。" * 3
    history = [
        "用户：帮我总结一下 2026 年 Q2 的销售数据报告的主要发现。" * 2,
        "助手：根据报告，Q2 收入同比增长 12%，主要由欧洲区贡献……" * 4,
    ] * 10  # 模拟长对话

    plan_context_budget(
        model="claude-sonnet-5",
        system=system,
        history_msgs=history,
        tool_defs_tokens=3_200,      # 15 个工具，粗估每个 ~200 token
        output_reserve=8_000,
    )

    print()
    # 中文 vs 英文 token 效率对比（为何不能用字符数做预算）
    zh_text = "大语言模型通过注意力机制理解上下文。" * 10
    en_text = "Large language models understand context via attention." * 10
    print(f"中文 {len(zh_text)} 字符 ≈ {estimate_tokens(zh_text)} tokens")
    print(f"英文 {len(en_text)} 字符 ≈ {estimate_tokens(en_text)} tokens")
    print("结论：同等信息量下中文 token 成本更高，预算必须按 token 而非字符做。")

    # 图片 token 速查：与官方 vision 文档表对照（1000x1000 → 1296）
    assert image_tokens(1000, 1000) == 1296, "与官方表不一致"
    assert image_tokens(200, 200) == 64, "与官方表不一致"
    print("\n图片 token 校验通过：200x200→64，1000x1000→1296（与官方表一致）")