#!/usr/bin/env python3
"""api_walkthrough.py — Messages / Batches / count_tokens 请求构建与校验

两种运行模式：
  1. 校验模式（默认，无 key 可跑）：构建请求并按官方约束做本地校验，
     打印将要发送的精确结构（绝不包含密钥）。
  2. 实发模式：设置环境变量 ANTHROPIC_API_KEY 后，用 urllib 真实调用
     POST /v1/messages（max_tokens=32 的小请求）与 count_tokens。

校验覆盖的官方约束（访问于 2026-08-30）:
  - Messages: model/max_tokens/messages 必填；请求体 <= 32MB
    https://platform.claude.com/docs/en/api/messages
  - Batches: custom_id 匹配 ^[a-zA-Z0-9_-]{1,64}$；单批 <= 100,000 请求或 256MB；
    不支持 stream / speed / store / cache_hint / max_tokens:0
    https://platform.claude.com/docs/en/build-with-claude/batch-processing
  - Prompt caching: cache_control 断点最多 4 个；ttl 只能 5m/1h
    https://platform.claude.com/docs/en/build-with-claude/prompt-caching

运行:
  python3 api_walkthrough.py           # 校验模式
  python3 api_walkthrough.py --live    # 实发模式（需要环境变量中的 key）
"""

from __future__ import annotations

import json
import os
import re
import sys
import urllib.request

API_BASE = "https://api.anthropic.com"
API_VERSION = "2023-06-01"
CUSTOM_ID_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")
BATCH_UNSUPPORTED = {"stream", "speed", "store", "previous_thread_event_id",
                     "cache_hint", "context_hint"}


def build_messages_request() -> dict:
    """带显式缓存断点的标准请求：稳定内容在前（可缓存），易变内容在后。"""
    return {
        "model": "claude-sonnet-5",
        "max_tokens": 1024,
        "system": [
            {
                "type": "text",
                "text": "你是企业知识库助手。回答必须引用来源文档段落。"
                        "（此处省略 2000 token 的检索上下文——稳定前缀，可缓存）",
                "cache_control": {"type": "ephemeral"},  # 显式断点：稳定前缀
            }
        ],
        "messages": [
            {"role": "user", "content": "Q3 欧洲区的收入增长是多少？（每轮变化的问题）"},
        ],
    }


def build_batch_request(items: list[tuple[str, dict]]) -> dict:
    return {"requests": [{"custom_id": cid, "params": params} for cid, params in items]}


def validate_messages_request(req: dict) -> list[str]:
    errs = []
    for field in ("model", "max_tokens", "messages"):
        if field not in req:
            errs.append(f"Messages: 缺必填字段 {field}")
    if req.get("max_tokens", 1) < 0:
        errs.append("Messages: max_tokens 不能为负（0 仅用于缓存预热）")
    if len(json.dumps(req)) > 32 * 1024 * 1024:
        errs.append("Messages: 超过 32MB 请求体上限（413）")
    n_breaks = json.dumps(req).count('"cache_control"')
    if n_breaks > 4:
        errs.append(f"Prompt caching: 断点 {n_breaks} 个 > 上限 4")
    for blk in req.get("system", []) if isinstance(req.get("system"), list) else []:
        ttl = (blk.get("cache_control") or {}).get("ttl", "5m")
        if ttl not in ("5m", "1h"):
            errs.append(f"Prompt caching: ttl 非法 {ttl}（只能 5m/1h）")
    return errs


def validate_batch_request(batch: dict) -> list[str]:
    errs = []
    reqs = batch.get("requests", [])
    if len(reqs) > 100_000:
        errs.append("Batches: 超过 100,000 请求上限")
    if len(json.dumps(batch)) > 256 * 1024 * 1024:
        errs.append("Batches: 超过 256MB 上限")
    for r in reqs:
        if not CUSTOM_ID_RE.match(r.get("custom_id", "")):
            errs.append(f"Batches: custom_id 非法 {r['custom_id'][:20]!r}")
        bad = BATCH_UNSUPPORTED & set(r.get("params", {}))
        if bad:
            errs.append(f"Batches: params 不支持 {sorted(bad)}")
        if r.get("params", {}).get("max_tokens") == 0:
            errs.append("Batches: max_tokens=0 不支持（缓存预热不能进批）")
    return errs


def live_call(req: dict) -> None:
    """实发模式：小请求走 Messages；密钥只从环境读取，绝不写入任何输出。"""
    key = os.environ.get("ANTHROPIC" + "_API_KEY", "")
    if not key:
        sys.exit("实发模式需要环境变量提供 key（见文件头注释）")
    req = dict(req, max_tokens=32)  # 演示用小请求
    body = json.dumps(req).encode()
    http_req = urllib.request.Request(
        f"{API_BASE}/v1/messages",
        data=body,
        headers={
            "content-type": "application/json",
            "anthropic-version": API_VERSION,
            "x-api-key": key,  # 从环境注入；本脚本不打印任何头部
        },
        method="POST",
    )
    with urllib.request.urlopen(http_req, timeout=60) as resp:
        data = json.loads(resp.read())
        print("stop_reason:", data["stop_reason"])
        print("usage:", data["usage"])
        print("request-id:", resp.headers.get("request-id"))


def main() -> None:
    live = "--live" in sys.argv
    msg_req = build_messages_request()

    if live:
        live_call(msg_req)
        return

    print("=== 校验模式（无 key）：Messages 请求 ===")
    print(json.dumps(msg_req, ensure_ascii=False, indent=2)[:800], "...\n")
    errs = validate_messages_request(msg_req)
    print("Messages 校验:", "通过" if not errs else errs)

    batch = build_batch_request([
        ("eval-001", {"model": "claude-haiku-4-5", "max_tokens": 512,
                      "messages": [{"role": "user", "content": "判断情感:好评"}]}),
        ("bad id!", {"model": "claude-haiku-4-5", "max_tokens": 0,
                     "messages": [{"role": "user", "content": "x"}]}),
    ])
    print("\n=== 校验模式：Batches 请求（故意含 2 处错误） ===")
    errs = validate_batch_request(batch)
    for e in errs:
        print("  [!]", e)
    print("（正确做法：先在 Messages API 上验证 params 形状，再入批）")

    # 自检
    good_batch = build_batch_request([("ok_1", {"model": "claude-haiku-4-5",
                                                "max_tokens": 1,
                                                "messages": [{"role": "user", "content": "hi"}]})])
    assert validate_messages_request(msg_req) == []
    assert len(validate_batch_request(batch)) == 2
    assert validate_batch_request(good_batch) == []
    print("\n自检通过：请求构建与校验逻辑符合官方约束")


if __name__ == "__main__":
    main()