# /// script
# requires-python = ">=3.10"
# dependencies = ["anthropic>=2.3,<3", "pillow>=10"]
# ///
"""Give Claude a task on this VM's desktop, with the computer use tool.

    uv run agent.py "In Google Chrome, open news.ycombinator.com and tell me the top story."

Claude looks at screenshots of display :1 and answers with clicks and
keystrokes, which xdotool carries out. Each step is printed as it happens.
Screenshots, zooms and waits run unasked; before any other action you
answer y (run it), n (refuse it; Claude is told) or a (run it and every
later one). --yes approves everything from the start: only for tasks and
pages you trust, since what is on the screen steers what Claude does.

Each run is kept in ~/agent-runs/<time>/: log.txt, with the task, every
action and Claude's replies, and screen.mp4, a recording of the desktop
(--no-record leaves it out).

The model comes from exe.dev's LLM integration (ANTHROPIC_BASE_URL and
ANTHROPIC_API_KEY default to it); set MODEL or --model to pick another.
From https://exe.dev/deploy/computer-use.
"""

from __future__ import annotations

import argparse
import base64
import io
import json
import os
import shutil
import signal
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

os.environ.setdefault("ANTHROPIC_BASE_URL", "https://llm.int.exe.xyz")
os.environ.setdefault("ANTHROPIC_API_KEY", "implicit")

from anthropic import Anthropic, AnthropicError  # noqa: E402
from anthropic.tools import ToolError  # noqa: E402
from anthropic.tools.computer import (  # noqa: E402
    BetaAbstractComputerToolset20260801,
    BetaComputerConfirmContext,
    BetaComputerCursorPositionResult,
    BetaScreenshotResult,
)
from PIL import ImageGrab  # noqa: E402

DISPLAY = os.environ.get("DISPLAY_NUM", ":1")
MAX_WIDTH, MAX_HEIGHT = 1920, 1200  # screenshots go to the model unscaled
SETTLE = 0.4  # seconds for the desktop to catch up after an action
# What Claude should know about this desktop.
SYSTEM = """You are using an Xfce desktop on an Ubuntu Linux VM, through the computer tool.
Google Chrome is open; a terminal and a file manager are in the dock at the bottom.
To go to a web address in Chrome, press ctrl+l first (a click in the address bar
may only place the cursor), then type the address and press Return.
After an action that changes the screen, take a screenshot before relying on it."""
UNASKED = {"screenshot", "zoom", "wait", "mouse_move", "cursor_position"}
WHEEL = {"up": 4, "down": 5, "left": 6, "right": 7}
# xdotool takes X keysym names, which are case-sensitive; these are the
# other spellings a model may use.
KEYS = {
    "enter": "Return", "return": "Return", "esc": "Escape", "escape": "Escape", "backspace": "BackSpace",
    "tab": "Tab", "space": "space", "delete": "Delete", "del": "Delete", "insert": "Insert", "home": "Home",
    "end": "End", "pageup": "Page_Up", "page_up": "Page_Up", "pagedown": "Page_Down", "page_down": "Page_Down",
    "up": "Up", "down": "Down", "left": "Left", "right": "Right", "cmd": "super", "win": "super", "meta": "super",
    "control": "ctrl",
}  # fmt: skip


def keysyms(chord: str) -> str:
    """ctrl+Enter -> ctrl+Return."""
    return "+".join(KEYS.get(k.lower(), k) for k in chord.split("+"))


class Log:
    """Prints to the terminal and appends to the run's log.txt."""

    def __init__(self, path: Path) -> None:
        self.file = path.open("a", encoding="utf-8")

    def __call__(self, line: str, echo: bool = True) -> None:
        if echo:
            print(line, flush=True)
        self.file.write(f"{datetime.now():%H:%M:%S} {line}\n")
        self.file.flush()


class Desktop(BetaAbstractComputerToolset20260801):
    """The computer toolset on an X display, through xdotool."""

    def __init__(self, **options) -> None:
        super().__init__(**options)
        self.width, self.height = ImageGrab.grab(xdisplay=DISPLAY).size
        if self.width > MAX_WIDTH or self.height > MAX_HEIGHT:
            raise RuntimeError(f"the screen is {self.width}x{self.height}; use {MAX_WIDTH}x{MAX_HEIGHT} or less")

    def _xdo(self, *args: object) -> None:
        env = {**os.environ, "DISPLAY": DISPLAY}
        done = subprocess.run(["xdotool", *map(str, args)], env=env, capture_output=True, text=True)
        if done.returncode != 0:
            raise ToolError(f"xdotool failed: {done.stderr.strip()}")
        time.sleep(SETTLE)

    def _at(self, coordinate) -> list[object]:
        if coordinate is None:
            return []
        x, y = coordinate
        if not (0 <= x < self.width and 0 <= y < self.height):
            raise ToolError(f"({x}, {y}) is outside the {self.width}x{self.height} screen")
        return ["mousemove", x, y]

    def _click(self, coordinate, modifiers, button: int, times: int = 1) -> None:
        held = keysyms(modifiers).split("+") if modifiers else []
        args = self._at(coordinate)
        for key in held:
            args += ["keydown", key]
        args += ["click", "--repeat", times, "--delay", 80, button]
        for key in reversed(held):
            args += ["keyup", key]
        self._xdo(*args)

    def screenshot(self, context, input) -> BetaScreenshotResult:
        return self._png(ImageGrab.grab(xdisplay=DISPLAY))

    def zoom(self, context, input) -> BetaScreenshotResult:
        x0, y0, x1, y1 = input.region
        if not (0 <= x0 < x1 <= self.width and 0 <= y0 < y1 <= self.height):
            raise ToolError(f"region {input.region} is not inside the {self.width}x{self.height} screen")
        image = ImageGrab.grab(xdisplay=DISPLAY).crop((x0, y0, x1, y1))
        scale = min(self.width / image.width, self.height / image.height)
        return self._png(image.resize((int(image.width * scale), int(image.height * scale))))

    def _png(self, image) -> BetaScreenshotResult:
        out = io.BytesIO()
        image.convert("RGB").save(out, "PNG")
        return BetaScreenshotResult(data=base64.b64encode(out.getvalue()).decode())

    def cursor_position(self, context, input) -> BetaComputerCursorPositionResult:
        env = {**os.environ, "DISPLAY": DISPLAY}
        out = subprocess.run(["xdotool", "getmouselocation", "--shell"], env=env, capture_output=True, text=True).stdout
        at = dict(line.split("=", 1) for line in out.split())
        return BetaComputerCursorPositionResult(x=int(at["X"]), y=int(at["Y"]))

    def mouse_move(self, context, input) -> None:
        self._xdo(*self._at(input.coordinate))

    def left_click(self, context, input) -> None:
        self._click(input.coordinate, input.text, 1)

    def right_click(self, context, input) -> None:
        self._click(input.coordinate, input.text, 3)

    def middle_click(self, context, input) -> None:
        self._click(input.coordinate, input.text, 2)

    def double_click(self, context, input) -> None:
        self._click(input.coordinate, input.text, 1, times=2)

    def triple_click(self, context, input) -> None:
        self._click(input.coordinate, input.text, 1, times=3)

    def left_click_drag(self, context, input) -> None:
        self._xdo(*self._at(input.start_coordinate), "mousedown", 1, *self._at(input.coordinate), "mouseup", 1)

    def scroll(self, context, input) -> None:
        self._click(input.coordinate, input.text, WHEEL[input.scroll_direction], times=input.scroll_amount)

    def key(self, context, input) -> None:
        self._xdo("key", "--repeat", input.repeat or 1, "--delay", 50, "--", *[keysyms(chord) for chord in input.text.split()])

    def type(self, context, input) -> None:
        self._xdo("type", "--delay", 25, "--", input.text)

    def wait(self, context, input) -> None:
        time.sleep(input.duration)


def describe(member: str, input) -> str:
    """One line for an action, such as: left_click (880, 118)."""
    args = input.model_dump(exclude_none=True)
    where = args.pop("coordinate", None)
    parts = [member] + ([f"({where[0]}, {where[1]})"] if where else [])
    parts += [json.dumps(v) if k == "text" else f"{k}={v}" for k, v in args.items()]
    return " ".join(parts)


def main() -> None:
    parser = argparse.ArgumentParser(description="Give Claude a task on this VM's desktop.")
    parser.add_argument("task", nargs="+", help="what Claude should do")
    parser.add_argument("--yes", action="store_true", help="run every action without asking")
    parser.add_argument("--no-record", action="store_true", help="don't record the screen")
    parser.add_argument("--model", default=os.environ.get("MODEL", "claude-sonnet-5-5"))
    parser.add_argument("--max-steps", type=int, default=50, help="stop after this many model turns")
    args = parser.parse_args()
    task = " ".join(args.task)

    run = Path.home() / "agent-runs" / datetime.now().strftime("%Y-%m-%d-%H%M%S")
    run.mkdir(parents=True)
    log = Log(run / "log.txt")
    approve_all = args.yes

    def confirm(context: BetaComputerConfirmContext) -> bool:
        nonlocal approve_all
        line = describe(context.member, context.input)
        if approve_all or context.member in UNASKED:
            log(f"  > {line}")
            return True
        try:
            answer = input(f"  > {line}   run it? [y/n/a] ").strip().lower()
        except EOFError:  # no one at a terminal to ask, as under another agent
            print("(no terminal to ask: refused; use --yes to allow)")
            answer = "n"
        approve_all = answer == "a"
        ok = answer in ("y", "a")
        log(f"  > {line}" + ("" if ok else "   (refused)"), echo=False)
        return ok

    try:
        desktop = Desktop(confirm=confirm)
        client = Anthropic()
    except (AnthropicError, OSError, RuntimeError) as error:
        sys.exit(f"error: {error}")

    recorder = None
    if not args.no_record and shutil.which("ffmpeg"):
        recorder = subprocess.Popen(
            ["ffmpeg", "-loglevel", "error", "-f", "x11grab", "-framerate", "5",
             "-video_size", f"{desktop.width}x{desktop.height}", "-i", DISPLAY,
             "-c:v", "libx264", "-preset", "veryfast", "-crf", "30", "-pix_fmt", "yuv420p",
             "-movflags", "+frag_keyframe+empty_moov", str(run / "screen.mp4")],
            stdin=subprocess.DEVNULL, stderr=(run / "ffmpeg.log").open("w"), start_new_session=True,
        )  # fmt: skip

    log(f"Task: {task}")
    log(f"Model: {args.model}. Run kept in {run}")
    tokens_in = tokens_out = 0
    try:
        runner = client.beta.messages.tool_runner(
            model=args.model,
            max_tokens=4096,
            max_iterations=args.max_steps,
            system=SYSTEM,
            tools=[desktop],
            messages=[{"role": "user", "content": task}],
        )
        for message in runner:
            tokens_in += message.usage.input_tokens + (message.usage.cache_read_input_tokens or 0)
            tokens_out += message.usage.output_tokens
            for block in message.content:
                if block.type == "text" and block.text.strip():
                    log(f"Claude: {block.text.strip()}")
    except KeyboardInterrupt:
        log("Stopped.")
    except AnthropicError as error:
        log(f"error: {error}")
    finally:
        desktop.close()
        if recorder:
            recorder.send_signal(signal.SIGINT)
            recorder.wait(timeout=10)
        log(f"Tokens: {tokens_in:,} in, {tokens_out:,} out.")
        video = run / "screen.mp4"
        if recorder and video.exists() and video.stat().st_size > 0:
            log(f"Recording: {video}")


if __name__ == "__main__":
    main()