#!/usr/bin/env python3
"""message_builder.py — Messages API 消息结构构建与校验（无需 API key，本地可跑）

演示：
1. system 是顶层参数（Messages API 没有 "system" 角色）
2. user/assistant 交替轮次；content 可以是 string 或内容块数组
3. 多模态块（image base64 / document）的合法结构
4. prefill（末尾 assistant 消息）会被服务端"续写"
5. 发送前本地校验（角色交替、块类型、必填字段）

来源: https://platform.claude.com/docs/en/api/messages（访问于 2026-08-30）
运行: python3 message_builder.py
"""

import json
import sys

VALID_INPUT_BLOCKS = {"text", "image", "document", "tool_use", "tool_result", "thinking"}


def text_block(t: str) -> dict:
    return {"type": "text", "text": t}


def image_block_b64(media_type: str, data_b64: str) -> dict:
    return {
        "type": "image",
        "source": {"type": "base64", "media_type": media_type, "data": data_b64},
    }


def build_request(system: str, turns: list[dict], *, model="claude-sonnet-5",
                  max_tokens=1024, prefill: str | None = None) -> dict:
    """把 (role, content) 轮次组装成合法请求体。prefill 追加末尾 assistant 消息。"""
    messages = []
    for t in turns:
        messages.append({"role": t["role"], "content": t["content"]})
    if prefill is not None:
        messages.append({"role": "assistant", "content": prefill})
    return {
        "model": model,
        "max_tokens": max_tokens,
        "system": system,          # 顶层参数，不是消息角色
        "messages": messages,
    }


def validate_request(req: dict) -> list[str]:
    """发送前的应用层校验（返回问题列表，空列表=通过）。"""
    problems = []
    msgs = req.get("messages", [])
    if not msgs:
        problems.append("messages 为空")
    for i, m in enumerate(msgs):
        if m["role"] not in ("user", "assistant"):
            problems.append(f"[{i}] 非法角色 {m['role']!r}（只允许 user/assistant）")
        c = m["content"]
        if isinstance(c, list):
            for b in c:
                if isinstance(b, dict) and b.get("type") not in VALID_INPUT_BLOCKS:
                    problems.append(f"[{i}] 未知块类型 {b.get('type')!r}")
    if "max_tokens" not in req:
        problems.append("缺 max_tokens（必填）")
    if "model" not in req:
        problems.append("缺 model（必填）")
    return problems


if __name__ == "__main__":
    system = "你是财务分析助手。输出必须以 JSON 开头。"

    req = build_request(
        system=system,
        turns=[
            {"role": "user", "content": "分析这份报表"},
            {"role": "user", "content": [text_block("重点看毛利率"), image_block_b64("image/png", "iVBORw0KGgo=")]},
            {"role": "assistant", "content": "毛利率同比 +3pt，主要来自原材料降价。"},
            {"role": "user", "content": "那净利率呢？"},
        ],
        prefill='{"summary": "',   # 约束响应从 JSON 续写
    )

    print("== 请求体（注意 system 在顶层；连续 user 会被服务端合并）==")
    print(json.dumps(req, ensure_ascii=False, indent=2))

    problems = validate_request(req)
    print("\n== 本地校验 ==", "通过" if not problems else problems)

    # 故意放一个坏角色，演示校验器
    bad = dict(req)
    bad["messages"] = [{"role": "system", "content": "hi"}]
    bad_problems = validate_request(bad)
    assert any("system" in p for p in bad_problems), "校验器应拒绝 system 角色"
    print("坏角色拦截测试: OK（Messages API 无 system 角色 -> 应用层就该拦下）")
    sys.exit(0)