"""MCP 工具暴露面 vs 直接 API 集成：上下文成本模拟器（无需 API key）。

不接入真实 LLM，用一个确定性的 token 估算（chars/4，粗估）模拟两种架构：
A) 每轮把所有已连接 MCP server 的全部工具定义注入模型上下文（naive host）
B) 渐进式发现 progressive discovery：只注入 1 个 search_tools 元工具 + 按需加载
   （对应官方 Client Best Practices 的 catalog/inspect/execute 三层模式）

官方数据点：上百工具时，全量定义可占 ~150,000 tokens；渐进式 ~2,000 tokens。
来源: https://modelcontextprotocol.io/docs/2026-07-28/develop/clients/client-best-practices
（访问于 2026-08-30）

运行：
    uv run --with mcp python 03_context_cost_simulator.py
"""

import json

CONTEXT_WINDOW = 200_000  # 以 200k token 上下文为例（示意值）
META_TOOL_DEF = {
    "name": "search_tools",
    "description": "Search available tools by query; returns names + one-line descriptions.",
    "inputSchema": {"type": "object", "properties": {"query": {"type": "string"}},
                    "required": ["query"]},
}


def make_tool_def(server: str, i: int) -> dict:
    """模拟第 i 个工具定义（名称、描述、含 6 个字段的 inputSchema）。"""
    props = {
        f"field_{j}": {"type": "string", "description": f"Field {j} of record {i} on {server}"}
        for j in range(6)
    }
    return {
        "name": f"{server}_operation_{i}",
        "description": f"Perform operation {i} against the {server} system (create/update/read).",
        "inputSchema": {"type": "object", "properties": props,
                        "required": list(props)[:3]},
    }


def est_tokens(obj) -> int:
    """粗略 token 估算：JSON 序列化后按 ~4 chars/token。"""
    return max(1, len(json.dumps(obj)) // 4)


def simulate(servers: dict[str, int], task_needs: int):
    """servers: {server_name: tool_count}; task_needs: 本次任务实际用到的工具数。"""
    all_defs = []
    for name, count in servers.items():
        all_defs += [make_tool_def(name, i) for i in range(count)]

    naive_tokens = sum(est_tokens(d) for d in all_defs)
    # 渐进式：元工具 + 任务实际用到的少数几个定义（inspect 层按需加载）
    progressive_tokens = est_tokens(META_TOOL_DEF) + sum(
        est_tokens(d) for d in all_defs[:task_needs])

    print(f"场景：连接 {len(servers)} 个 MCP server，共 {len(all_defs)} 个工具；"
          f"本次任务只需 {task_needs} 个\n")
    print(f"上下文窗口（示意）        : {CONTEXT_WINDOW:>9,} tokens")
    print(f"A. 全量注入工具定义        : {naive_tokens:>9,} tokens "
          f"({naive_tokens / CONTEXT_WINDOW:.0%} 窗口)")
    print(f"B. progressive discovery  : {progressive_tokens:>9,} tokens "
          f"({progressive_tokens / CONTEXT_WINDOW:.0%} 窗口)")
    print(f"节省                      : {(1 - progressive_tokens / naive_tokens):.1%}\n")

    print("判读：")
    if naive_tokens / CONTEXT_WINDOW > 0.05:
        print("- 工具定义已超过官方建议的切换阈值（上下文窗口的 1%–5%），"
              "应启用 progressive discovery。")
    else:
        print("- 工具定义占比很小，全量注入完全可接受，无需过度工程。")
    print("- 这解释了「何时用 MCP vs 直接 API」的一个关键维度：工具是 *注入上下文的*，"
          "普通 REST 集成的 SDK 不是。")


if __name__ == "__main__":
    # 小规模：全量注入毫无压力
    simulate({"crm": 5, "erp": 4, "wiki": 3}, task_needs=2)
    print("=" * 72)
    # 规模化：上百工具时全量注入吃掉大半个窗口
    simulate({"crm": 60, "erp": 50, "helpdesk": 40, "wiki": 30, "monitoring": 30},
             task_needs=3)