#!/usr/bin/env python3
"""成本 / 延迟 / 可靠性权衡估算器(无需 API key)。

输入一个 workload 规格(每日请求数、稳定前缀 tokens、动态后缀 tokens、输出 tokens、
缓存命中率), 输出:
  1. 各模型档位的月成本与相对延迟(基于官方 comparative latency 定性排序映射为启发式系数)
  2. 开/关 prompt caching 的成本对比(写入 1.25x/2x, 读取 0.1x)
  3. 实时 API vs Message Batches(半价)对比
  4. 一条基于规则的推荐(实时交互不进 batch; 离线批量进 batch; 稳定前缀>1k tokens 上缓存)

价格数据截至 2026-08-30, 来源:
  - 模型价: https://platform.claude.com/docs/en/models/overview
  - 缓存计价倍率(5m 写 1.25x / 1h 写 2x / 读 0.1x)与 Batches 半价:
    https://platform.claude.com/docs/en/build-with-claude/prompt-caching
    https://platform.claude.com/docs/en/build-with-claude/batch-processing
延迟系数为启发式假设(官方仅给 slower/moderate/fast/fastest 定性排序), 已在表中注明。

运行: python3 02_cost_latency_estimator.py
"""

from __future__ import annotations

from dataclasses import dataclass

# ------------------------------------------------------------------ 价格表
# 截至 2026-08-30; 单位: USD / 百万 tokens
MODELS = {
    # id: (名称, 输入价, 输出价, 延迟系数[启发式: 官方 slower->1.6 / moderate->1.3 /
    #      fast->1.0 / fastest->0.6], 官方定位)
    "claude-fable-5":  ("Claude Fable 5", 10.0, 50.0, 1.6, "长时 agent 的下一代智能(较慢)"),
    "claude-opus-5":   ("Claude Opus 5",   5.0, 25.0, 1.3, "复杂 agentic 编码与企业工作"),
    "claude-sonnet-5": ("Claude Sonnet 5", 2.0, 10.0, 1.0, "速度与智能的最佳组合"),
    "claude-haiku-4-5-20251001": ("Claude Haiku 4.5", 1.0, 5.0, 0.6, "最快、近前沿智能"),
}

CACHE_WRITE_5M = 1.25   # 5 分钟缓存写入 = 基础输入价 x1.25 (官方)
CACHE_WRITE_1H = 2.0    # 1 小时缓存写入 = x2.0 (官方)
CACHE_READ = 0.1        # 缓存读取 = x0.1 (官方)
BATCH_DISCOUNT = 0.5    # Message Batches 全量半价 (官方)


@dataclass
class Workload:
    name: str
    requests_per_day: int
    stable_prefix_tokens: int   # 跨请求稳定(可缓存): system + 工具定义 + 文档前缀
    dynamic_tokens: int         # 每次变化(当次问题/时间戳等)
    output_tokens: int
    cache_hit_rate: float       # 开缓存后的命中率假设
    interactive: bool           # 用户是否在等结果(决定能否进 batch)
    reliability_need: str       # low / medium / high — 用于档位过滤


# ------------------------------------------------------------------ 计算器
def monthly_cost(model_id: str, w: Workload, *, use_cache: bool,
                 use_batch: bool) -> float:
    name, p_in, p_out, _, _ = MODELS[model_id]
    n = w.requests_per_day * 30
    input_cost = n * (w.stable_prefix_tokens + w.dynamic_tokens) * p_in / 1e6
    output_cost = n * w.output_tokens * p_out / 1e6
    cost = input_cost + output_cost

    if use_cache:
        hit = w.cache_hit_rate
        miss = 1.0 - hit
        # 未命中部分: 全额重写缓存(1.25x); 命中部分: 0.1x 读取
        stable = w.stable_prefix_tokens
        input_cost = n * (stable * miss * CACHE_WRITE_5M * p_in
                          + stable * hit * CACHE_READ * p_in
                          + w.dynamic_tokens * p_in) / 1e6
        cost = input_cost + output_cost

    if use_batch:
        cost *= BATCH_DISCOUNT
    return cost


def latency_score(model_id: str) -> float:
    """相对延迟系数(启发式): Sonnet=1.0 基准。官方仅给定性排序。"""
    return MODELS[model_id][3]


def analyze(w: Workload) -> None:
    print(f"\n===== Workload: {w.name} =====")
    print(f"  {w.requests_per_day:,} 请求/日; 前缀 {w.stable_prefix_tokens:,} + "
          f"动态 {w.dynamic_tokens:,} + 输出 {w.output_tokens:,} tokens; "
          f"交互式={w.interactive}; 可靠性要求={w.reliability_need}")

    print(f"\n  {'模型':<30}{'无缓存':>12}{'缓存@' + f'{int(w.cache_hit_rate*100)}%命中':>14}"
          f"{'批处理(半价)':>14}{'相对延迟':>10}")
    rows = []
    for mid in MODELS:
        base = monthly_cost(mid, w, use_cache=False, use_batch=False)
        cached = monthly_cost(mid, w, use_cache=True, use_batch=False)
        batched = monthly_cost(mid, w, use_cache=False,
                               use_batch=not w.interactive)
        rows.append((mid, base, cached, batched))
        print(f"  {MODELS[mid][0]:<30}{base:>11,.0f}${cached:>13,.0f}$"
              f"{batched:>13,.0f}${latency_score(mid):>9.1f}x")

    # ---- 推荐逻辑(规则显式可审计) ------------------------------------
    candidates = list(MODELS)
    if w.reliability_need == "high":
        candidates = ["claude-opus-5", "claude-fable-5"]
    elif w.reliability_need == "medium":
        candidates = ["claude-sonnet-5", "claude-opus-5"]  # 质量与延迟都在意
    else:  # low
        candidates = ["claude-haiku-4-5-20251001", "claude-sonnet-5"]

    rec_model = min(candidates,
                    key=lambda m: monthly_cost(m, w, use_cache=True,
                                               use_batch=not w.interactive))
    rec_cost = monthly_cost(rec_model, w, use_cache=True,
                            use_batch=not w.interactive)
    print("\n  推荐决策链:")
    print(f"    1. 可靠性要求 {w.reliability_need} -> 候选档位 "
          f"{[MODELS[m][0] for m in candidates]}")
    tag = "交互式(实时 API, 不可 batch)" if w.interactive else "离线(进 Batches, 半价)"
    print(f"    2. 交互性: {tag}")
    if w.stable_prefix_tokens >= 1000:
        print(f"    3. 稳定前缀 {w.stable_prefix_tokens:,} tokens >= 1k -> "
              f"开 prompt caching; cache_control 放在稳定前缀末块")
    else:
        print("    3. 稳定前缀 < 1k tokens -> 缓存收益有限, 不强制开缓存")
    print(f"    => 推荐: {MODELS[rec_model][0]} + "
          f"{'缓存' if w.stable_prefix_tokens >= 1000 else '无缓存'} + "
          f"{'Batches' if not w.interactive else '实时 API'}"
          f" ≈ ${rec_cost:,.0f}/月 (估算, 价格截至 2026-08-30)")


def main() -> None:
    # 自检: 计价倍率来自官方文档
    assert CACHE_WRITE_5M == 1.25 and CACHE_WRITE_1H == 2.0 and CACHE_READ == 0.1
    assert BATCH_DISCOUNT == 0.5

    # 场景 A: 客服聊天(交互式, 中可靠性) — 缓存收益典型场景
    analyze(Workload(
        name="客服助手(交互式)",
        requests_per_day=20_000,
        stable_prefix_tokens=6_000,   # system + 产品文档 + 工具定义
        dynamic_tokens=300,
        output_tokens=500,
        cache_hit_rate=0.75,
        interactive=True,
        reliability_need="medium",
    ))

    # 场景 B: 夜间批量评估(离线, 低交互)
    analyze(Workload(
        name="夜间批量内容审核(离线)",
        requests_per_day=50_000,
        stable_prefix_tokens=2_000,
        dynamic_tokens=800,
        output_tokens=200,
        cache_hit_rate=0.8,
        interactive=False,
        reliability_need="low",
    ))

    # 场景 C: 高价值合同审查(高可靠性)
    analyze(Workload(
        name="合同审查(高可靠性)",
        requests_per_day=500,
        stable_prefix_tokens=30_000,  # 长文档前缀
        dynamic_tokens=2_000,
        output_tokens=3_000,
        cache_hit_rate=0.6,
        interactive=True,
        reliability_need="high",
    ))

    # 缓存计价手工核算(可手算验证):
    # Sonnet 5, 1 次请求, 前缀 6k 全命中: 6000*0.2$/MTok = 0.0012$ 读取价
    # 未命中写入: 6000*2.5$ = 0.015$; 全额无缓存: 6000*2.0$ = 0.012$
    name, p_in, p_out, _, _ = MODELS["claude-sonnet-5"]
    hit = 6000 * CACHE_READ * p_in / 1e6
    write = 6000 * CACHE_WRITE_5M * p_in / 1e6
    full = 6000 * p_in / 1e6
    print("\n===== 手工核算(缓存读取/写入/全额, Sonnet 5, 6k tokens) =====")
    print(f"  命中读取 ${hit:.4f} vs 首次写入 ${write:.4f} vs 无缓存全额 ${full:.4f}")
    print(f"  => 命中一次省 {(1 - hit / full) * 100:.0f}% 输入成本; "
          f"写入比全额贵 25%; 命中 1 次即回本(0.1x vs 1.25x)")
    assert hit < full < write

    print("\n结论: 成本/延迟/可靠性三角没有全赢解 — 交互式买延迟(batch 不可用), "
          "离线买成本(半价), 高价值买可靠性(高档位+校验)。")


if __name__ == "__main__":
    main()
