"""Phase B: the same Claude agent does the same real tasks through each browser tool.

Usage: phase_b.py <tool> [task_id ...]   tool = playwright-mcp | devtools-mcp | agent-browser-mcp | agent-browser | agent-browser-agentsmd | agent-browser-noskill | agent-browser-lite
Each task gets a fresh browser. Records every turn's token usage, every tool call, time and the final answer.
"""
# ab_skill_core.md is the output of `agent-browser skills get core` (agent-browser 0.39.3)
import asyncio, base64, json, os, pathlib, shlex, subprocess, sys, time
import anthropic
from mcp import ClientSession
from mcp.client.stdio import stdio_client, StdioServerParameters

MODEL, EFFORT, MAX_TURNS = "claude-opus-5-5", "medium", 30
client = anthropic.Anthropic(max_retries=6)  # reads ANTHROPIC_API_KEY
TOOL = sys.argv[1]
TASKS = [t for t in json.load(open("tasks.json")) if len(sys.argv) < 3 or t["id"] in sys.argv[2:]]
OUT = pathlib.Path(os.environ.get("RUN", "phase_b")) / TOOL
OUT.mkdir(parents=True, exist_ok=True)

SYSTEM = ("You control a web browser through the tools you have. Complete the user's task in the browser. "
          "When you have the answer, reply with a short final answer and stop calling tools.")
SERVERS = {"playwright-mcp": ["@playwright/mcp", "--headless", "--isolated"],
           "agent-browser-mcp": ["agent-browser", "mcp"],  # default tool profile: core
           "devtools-mcp": ["chrome-devtools-mcp", "--headless", "--isolated", "--no-usage-statistics", "--no-performance-crux"]}
# agent-browser-lite: 3 rules from the core skill's own guidance, plus its output cap, instead of the 14K-token skill
LITE = ("agent-browser tips:\n"
        "- Run `snapshot -i` to see interactive elements only, and add `-s \"<css selector>\"` to scope it. Run it again after the page changes.\n"
        "- Use the @eN refs from the snapshot in commands like `click @e3`, `fill @e3 \"text\"` and `get text @e3`.\n"
        "- Inline `eval \"...\"` works only for simple expressions. Put the whole script in one double-quoted argument.")
# agent-browser-agentsmd: the block agent-browser's README suggests adding to AGENTS.md / CLAUDE.md, word for word
AGENTS_MD = """## Browser Automation

Use `agent-browser` for web automation. Run `agent-browser --help` for all commands.

Core workflow:

1. `agent-browser open <url>` - Navigate to page
2. `agent-browser snapshot -i` - Get interactive elements with refs (@e1, @e2)
3. `agent-browser click @e1` / `fill @e2 "text"` - Interact using refs
4. Re-snapshot after page changes"""
AB_TOOL = {"name": "agent_browser",
           "description": "Run one agent-browser CLI command. Pass everything after `agent-browser`, for example `open https://example.com` or `snapshot -i`.",
           "input_schema": {"type": "object", "properties": {"args": {"type": "string"}}, "required": ["args"]}}


def mcp_content(result):
    """MCP tool result -> Claude tool_result content (text and images, as the tool returned them)."""
    blocks = []
    for c in result.content:
        if c.type == "text":
            blocks.append({"type": "text", "text": c.text or " "})
        elif c.type == "image":
            blocks.append({"type": "image", "source": {"type": "base64", "media_type": c.mime_type, "data": c.data}})
    return blocks or [{"type": "text", "text": "(no output)"}]


def run_agent_browser(args, session, env=None):
    try:
        argv = shlex.split(args)
    except ValueError as e:
        return [{"type": "text", "text": f"error parsing arguments: {e}"}], True
    p = subprocess.run(["npx", "agent-browser", "--session", session] + argv, capture_output=True, text=True, timeout=180, env=env)
    text = (p.stdout + ("\n" + p.stderr if p.returncode else "")).strip() or "(no output)"
    return [{"type": "text", "text": text}], p.returncode != 0


async def run_task(task, call_tool, tools, system):
    messages = [{"role": "user", "content": task["prompt"]}]
    log = {"task": task["id"], "tool": TOOL, "model": MODEL, "effort": EFFORT, "turns": [], "tool_calls": []}
    t0 = time.time()
    if os.environ.get("PREFIX_CACHE"):  # cache tools + system once, reused across tasks, like a real agent harness
        system = [{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}]
    for turn in range(MAX_TURNS):
        resp = client.messages.create(model=MODEL, max_tokens=16000, system=system, tools=tools, messages=messages,
                                      output_config={"effort": EFFORT}, cache_control={"type": "ephemeral"})
        u = resp.usage
        log["turns"].append({"input": u.input_tokens, "cache_read": u.cache_read_input_tokens or 0,
                             "cache_write": u.cache_creation_input_tokens or 0, "output": u.output_tokens, "stop": resp.stop_reason})
        messages.append({"role": "assistant", "content": resp.content})
        if resp.stop_reason != "tool_use":
            log["final"] = "\n".join(b.text for b in resp.content if b.type == "text")
            break
        results = []
        for b in resp.content:
            if b.type != "tool_use":
                continue
            ts = time.time()
            try:
                content, is_error = await call_tool(b.name, b.input)
            except Exception as e:
                content, is_error = [{"type": "text", "text": f"tool error: {type(e).__name__}: {e}"}], True
            chars = sum(len(c.get("text", "")) for c in content)
            log["tool_calls"].append({"name": b.name, "input": b.input, "chars": chars, "images": sum(c["type"] == "image" for c in content),
                                      "error": is_error, "secs": round(time.time() - ts, 2),
                                      "text": "\n".join(c.get("text", "") for c in content)[:20000]})
            results.append({"type": "tool_result", "tool_use_id": b.id, "content": content, "is_error": is_error})
        messages.append({"role": "user", "content": results})
    else:
        log["final"] = None
        log["hit_turn_limit"] = True
    log["secs"] = round(time.time() - t0, 1)
    log["totals"] = {k: sum(t[k] for t in log["turns"]) for k in ("input", "cache_read", "cache_write", "output")}
    log["totals"]["input_all"] = log["totals"]["input"] + log["totals"]["cache_read"] + log["totals"]["cache_write"]
    return log


async def main():
    for task in TASKS:
        path = OUT / f"{task['id']}.json"
        if path.exists():
            continue
        if TOOL.startswith("agent-browser") and TOOL != "agent-browser-mcp":
            session = f"bench-{TOOL}-{task['id']}"
            system = {"agent-browser": SYSTEM + "\n\n" + open("ab_skill_core.md").read(),
                      "agent-browser-noskill": SYSTEM, "agent-browser-lite": SYSTEM + "\n\n" + LITE,
                      "agent-browser-agentsmd": SYSTEM + "\n\n" + AGENTS_MD}[TOOL]
            env = dict(os.environ, AGENT_BROWSER_MAX_OUTPUT="20000") if TOOL == "agent-browser-lite" else None

            async def call(name, inp, session=session, env=env):
                return run_agent_browser(inp.get("args", ""), session, env)

            log = await run_task(task, call, [AB_TOOL], system)
            subprocess.run(["npx", "agent-browser", "--session", session, "close"], capture_output=True, timeout=60)
        else:
            env = dict(os.environ, AGENT_BROWSER_SESSION=f"bench-mcp-{task['id']}") if TOOL == "agent-browser-mcp" else None
            async with stdio_client(StdioServerParameters(command="npx", args=SERVERS[TOOL], env=env)) as (r, w):
                async with ClientSession(r, w) as s:
                    await s.initialize()
                    tools = [{"name": t.name, "description": t.description or "", "input_schema": t.input_schema}
                             for t in (await s.list_tools()).tools]

                    async def call(name, inp, s=s):
                        res = await asyncio.wait_for(s.call_tool(name, inp), timeout=180)
                        return mcp_content(res), bool(res.is_error)

                    log = await run_task(task, call, tools, SYSTEM)
            if TOOL == "agent-browser-mcp":
                subprocess.run(["npx", "agent-browser", "--session", f"bench-mcp-{task['id']}", "close"], capture_output=True, timeout=60)
        json.dump(log, open(path, "w"), indent=1)
        print(f"{TOOL:<15} {task['id']:<26} turns={len(log['turns']):>2} calls={len(log['tool_calls']):>2} "
              f"in={log['totals']['input_all']:>8} out={log['totals']['output']:>6} {log['secs']:>6}s | {str(log.get('final'))[:90]!r}", flush=True)


asyncio.run(main())