"""Fixed context cost of each setup, in Claude tokens (messages.count_tokens, claude-opus-5-5). Saves overhead.json.
Needs anthropic.env with ANTHROPIC_API_KEY, and rea-agents 6.8.0 installed here (run with Node 24)."""
import asyncio, json, os, anthropic
from mcp import ClientSession
from mcp.client.stdio import stdio_client, StdioServerParameters

KEY = dict(l.strip().split("=", 1) for l in open("anthropic.env") if "=" in l)["ANTHROPIC_API_KEY"]
client = anthropic.Anthropic(api_key=KEY)
MSG = [{"role": "user", "content": "Find the API this page loads its data from."}]
SKILL = open("node_modules/rea-agents/skills/reverse-engineer-anything/SKILL.md").read()


async def rea_tools(env):
    params = StdioServerParameters(command="npx", args=["rea", "mcp"], env={**os.environ, **env})
    async with stdio_client(params) as (r, w):
        async with ClientSession(r, w) as s:
            await s.initialize()
            return [{"name": t.name, "description": t.description or "", "input_schema": t.input_schema} for t in (await s.list_tools()).tools]


def count(tools=None, system=None):
    kw = {k: v for k, v in (("tools", tools), ("system", system)) if v}
    base = client.messages.count_tokens(model="claude-opus-5-5", messages=MSG).input_tokens
    return client.messages.count_tokens(model="claude-opus-5-5", messages=MSG, **kw).input_tokens - base


async def main():
    compact, full = await rea_tools({}), await rea_tools({"REA_MCP_INPUT_SCHEMA_PROFILE": "full"})
    web = [t for t in compact if "web" in t["name"] or "browser" in t["name"]]
    pw = json.load(open("tools_playwright-mcp.json"))
    out = {"rea 6.8.0, default (compact)": [len(compact), count(compact)],
           "rea 6.8.0, full input schemas": [len(full), count(full)],
           "rea 6.8.0, web and browser tools": [len(web), count(web)],
           "rea skill (SKILL.md)": [0, count(system=SKILL)],
           "Playwright MCP 0.0.83": [len(pw), count(pw)]}
    json.dump(out, open("overhead.json", "w"), indent=1)

asyncio.run(main())