#!/usr/bin/env python3
"""Computer use agent loop skeleton (client-side execution) for the Anthropic Messages API.

STATUS (2026-09-18):
  * PARTS THAT RAN LIVE: the request shape for `computer_20251124` + header `computer-use-2025-11-24` on claude-sonnet-4-6
    (h1: Claude returned tool_use {"action":"screenshot"}) and the follow-up turn returning a 1x1 PNG screenshot as an
    image tool_result (h2: Claude described the screen). `computer_toolset_20260801` on claude-sonnet-5 (h4) returned
    tool_use {name:"screenshot", toolset_name:"computer", input:{}}. `computer_20250124` on haiku 4.5 (h3) also ran.
  * PARTS THAT DID NOT RUN: no real display was driven. `execute_action` below is a stub that returns a placeholder PNG for
    screenshot/zoom and "OK" for other actions. Wire it to xdotool / pyautogui / a VM in a sandbox before real use.

Run (2 requests, ≈ $0.013 on sonnet 4.6):  .venv/bin/python examples/anthropic/computer-use/computer_use_loop.py [--toolset]
  --toolset  use computer_toolset_20260801 (GA, needs claude-sonnet-5 / opus-5 / opus-4-8 / fable-5*) instead of computer_20251124 (beta).
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT))
from scripts.live import anthropic_request  # noqa: E402

PLACEHOLDER_PNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
NOT_EXECUTED = "Not executed: an earlier computer action in this turn failed."  # exact halt text (toolset batch actions)
USE_TOOLSET = "--toolset" in sys.argv
MODEL = "claude-sonnet-5" if USE_TOOLSET else "claude-sonnet-4-6"
BETA = None if USE_TOOLSET else "computer-use-2025-11-24"
TOOLS = [{"type": "computer_toolset_20260801", "configs": {"zoom": {"enabled": False}}}] if USE_TOOLSET else [
    {"type": "computer_20251124", "name": "computer", "display_width_px": 1024, "display_height_px": 768, "display_number": 1}]
SCALE = 1.0  # screenshot px -> real display px (coordinates come back in screenshot space; multiply by 1/SCALE before clicking)


def execute_action(action: str, params: dict):
    """STUB — replace with real display automation. Returns str or a list of content blocks."""
    if action in ("screenshot", "zoom"):
        return [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PLACEHOLDER_PNG}}]
    if action == "cursor_position":
        return "X=512, Y=384"
    if action in ("left_click", "right_click", "middle_click", "double_click", "triple_click", "mouse_move", "left_click_drag", "scroll"):
        coord = params.get("coordinate") or params.get("start_coordinate")
        if coord and not (0 <= coord[0] < 1024 and 0 <= coord[1] < 768):
            raise ValueError("Coordinates out of bounds")
        return "OK"  # would call: click(coord[0] / SCALE, coord[1] / SCALE)
    if action in ("type", "key", "hold_key", "wait", "left_mouse_down", "left_mouse_up"):
        return "OK"
    raise ValueError(f"Unknown or unimplemented action: {action}")


def process_tool_calls(content: list[dict]) -> list[dict]:
    """Run every computer action in order; after the first failure, skip the rest with the halt text (batch actions)."""
    results, failed = [], False
    for block in content:
        if block.get("type") != "tool_use":
            continue
        is_member = block.get("toolset_name") == "computer"
        if not (is_member or block["name"] == "computer"):
            continue  # other tools (bash, text editor, custom) would be dispatched here
        action = block["name"] if is_member else block["input"].get("action")
        params = block["input"]
        res: dict = {"type": "tool_result", "tool_use_id": block["id"]}
        if is_member:
            res["toolset_name"] = "computer"  # mandatory echo for toolset members
        if failed:
            res.update(content=NOT_EXECUTED, is_error=True)
        else:
            try:
                res["content"] = execute_action(action, params)
            except Exception as err:  # noqa: BLE001
                res.update(content=f"Error: {err}", is_error=True)
                failed = True
        results.append(res)
    return results


def sampling_loop(prompt: str, max_iterations: int = 2, max_tokens: int = 100) -> list[dict]:
    messages = [{"role": "user", "content": prompt}]
    for i in range(max_iterations):
        st, resp, _ = anthropic_request("POST", "/v1/messages", {"model": MODEL, "max_tokens": max_tokens, "tools": TOOLS, "messages": messages},
                                        beta=BETA, note=f"computer_use_loop.py iter {i} {TOOLS[0]['type']}")
        if st != 200:
            raise RuntimeError(json.dumps(resp)[:300])
        messages.append({"role": "assistant", "content": resp["content"]})
        print(f"[{i}] stop_reason={resp['stop_reason']}", [(b.get("name"), b.get("input")) for b in resp["content"] if b["type"] == "tool_use"])
        results = process_tool_calls(resp["content"])
        if not results or resp["stop_reason"] != "tool_use":
            break
        messages.append({"role": "user", "content": results})
        max_tokens = 50
    return messages


if __name__ == "__main__":
    msgs = sampling_loop("Take a screenshot.")
    print(json.dumps(msgs[-1], indent=1)[:600])