#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["openai-agents>=0.24,<0.25", "exedev>=0.4,<0.5"]
# ///
"""An OpenAI Agents SDK sandbox agent, with an exe.dev VM for each session.

  agent.py "a task"              start a session on a new VM, and do the task
  agent.py -s NAME "more"        continue session NAME, on its VM
  agent.py --ls                 list sessions
  agent.py --rm NAME            delete session NAME and its VM

A session is named after its first VM. Its conversation and sandbox state
are kept in sessions/ beside this file, so a later run picks up where the
last one stopped, with the same files and packages.

Settings, from the environment or .env beside this file:
  EXE_API         the exe-dev integration that makes session VMs, attached to
                  this VM; default https://openai-agents.int.exe.xyz
  EXE_TOKEN       an exe.dev API key to use instead, for new, ls, ssh, and rm
  MODEL           default gpt-6.1-sol
  OPENAI_API_KEY  use OpenAI directly. Without it, models come from exe.dev's
                  LLM gateway, https://llm.int.exe.xyz, which only exe.dev VMs
                  reach.

https://exe.dev/deploy/openai-agents-sdk
"""

import argparse
import asyncio
import json
import os
import sys
from pathlib import Path

import exedev
from agents import (Runner, SQLiteSession, set_default_openai_client,
                    set_tracing_disabled)
from agents.run import RunConfig
from agents.sandbox import SandboxAgent, SandboxRunConfig
from agents.sandbox.capabilities import Capabilities
from openai import AsyncOpenAI

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from exe_sandbox import ExeSandboxClient, ExeSandboxClientOptions  # noqa: E402

SESSIONS = HERE / "sessions"

INSTRUCTIONS = """\
You work on your own exe.dev VM: Ubuntu 24.04, with passwordless sudo, Docker,
Python, uv, Node.js, and Go. Your workspace is /workspace. Run `hostname` for
the VM's name. Install what you need.

A command returns only when its output closes, so don't start servers with &.
Run them under systemd instead:
`sudo systemd-run --unit=<name> --uid=exedev --working-directory=$PWD <command>`,
listening on 0.0.0.0 on a port from 3000 to 9999. A server on port <port> is
at https://<hostname>.exe.xyz:<port>/ for the user, who signs in to exe.dev
to see it. Give them that link.
"""


def load_env():
    env = HERE / ".env"
    if env.exists():
        for line in env.read_text().splitlines():
            name, eq, value = line.strip().partition("=")
            if eq and not name.startswith("#"):
                os.environ.setdefault(name.strip(), value.strip())


def use_model_source():
    if os.environ.get("OPENAI_API_KEY"):
        return  # OpenAI, with traces in its dashboard
    # exe.dev's LLM gateway; it adds the credentials, so the key is a placeholder.
    set_default_openai_client(
        AsyncOpenAI(base_url="https://llm.int.exe.xyz/v1", api_key="exe.dev"),
        use_for_tracing=False)
    set_tracing_disabled(True)  # traces go to OpenAI, which needs an OpenAI key


def print_item(item):
    if item.type == "tool_call_item":
        raw = item.raw_item
        get = raw.get if isinstance(raw, dict) else lambda k, d=None: getattr(raw, k, d)
        args = get("arguments") or get("action") or get("input") or ""
        if isinstance(args, str):
            try:
                args = json.loads(args)
            except ValueError:
                pass
        if isinstance(args, dict):
            args = args.get("cmd") or args.get("command") or args.get("commands") or args
        print(f"  {item.tool_name or get('type')}: {str(args)[:200]}", file=sys.stderr)


async def run(client, name, prompt):
    state_file = SESSIONS / f"{name}.json" if name else None
    if name and not state_file.exists():
        sys.exit(f"agent: no session {name}; see --ls")
    if state_file:
        state = client.deserialize_session_state(json.loads(state_file.read_text()))
        sandbox = await client.resume(state)
    else:
        sandbox = await client.create(options=ExeSandboxClientOptions(
            on_exit="keep", comment=f"agents-sdk: {prompt[:100]}"))
        name = sandbox.state.vm_name
        state_file = SESSIONS / f"{name}.json"
        print(f"session {name}, on VM {name}", file=sys.stderr)

    def save():
        SESSIONS.mkdir(exist_ok=True)
        state_file.write_text(json.dumps(client.serialize_session_state(sandbox.state)))

    save()  # before the run, so a failed run can be continued
    agent = SandboxAgent(
        name="Sandbox agent", model=os.environ.get("MODEL", "gpt-6.1-sol"),
        instructions=INSTRUCTIONS, capabilities=Capabilities.default())
    async with sandbox:
        result = Runner.run_streamed(
            agent, prompt, max_turns=100,
            session=SQLiteSession(name, SESSIONS / "history.db"),
            run_config=RunConfig(sandbox=SandboxRunConfig(session=sandbox),
                                 workflow_name="exe.dev sandbox agent"))
        async for event in result.stream_events():
            if event.type == "run_item_stream_event":
                print_item(event.item)
        print(result.final_output)
    save()
    print(f"\ncontinue: {sys.argv[0]} -s {name} \"...\"", file=sys.stderr)


async def remove(client, name):
    state_file = SESSIONS / f"{name}.json"
    if not state_file.exists():
        sys.exit(f"agent: no session {name}")
    state = client.deserialize_session_state(json.loads(state_file.read_text()))
    try:
        await (await client.exe.vm(state.vm_name)).delete()
        print(f"deleted VM {state.vm_name}")
    except Exception as e:  # already gone, most likely
        print(f"VM {state.vm_name}: {e}", file=sys.stderr)
    await SQLiteSession(name, SESSIONS / "history.db").clear_session()
    state_file.unlink()


async def main():
    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    p.add_argument("prompt", nargs="?")
    p.add_argument("-s", "--session", help="continue this session")
    p.add_argument("--ls", action="store_true", help="list sessions")
    p.add_argument("--rm", metavar="NAME", help="delete a session and its VM")
    a = p.parse_args()
    load_env()
    if a.ls:
        for f in sorted(SESSIONS.glob("*.json")):
            vm = json.loads(f.read_text())["vm_name"]
            print(f.stem, "" if vm == f.stem else f"(now on VM {vm})")
        return
    use_model_source()
    if os.environ.get("EXE_TOKEN"):
        client = ExeSandboxClient(os.environ["EXE_TOKEN"])
    else:
        client = ExeSandboxClient(api=os.environ.get("EXE_API", "https://openai-agents.int.exe.xyz"))
    try:
        if a.rm:
            await remove(client, a.rm)
        elif a.prompt:
            await run(client, a.session, a.prompt)
        else:
            p.print_usage()
    except exedev.ExeError as e:
        detail = (getattr(e, "body", "") or "").strip()
        sys.exit(f"agent: {e}" + (f": {detail}" if detail and detail not in str(e) else ""))
    finally:
        await client.aclose()


if __name__ == "__main__":
    asyncio.run(main())