#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["openai>=3.40,<4", "exedev>=0.4,<0.5"]
# ///
"""Run an OpenAI Agents API task, with its tools on an exe.dev VM of its own.

  agents_api.py "a task"            start a session on a new VM, and do the task
  agents_api.py -s SESSION "more"   continue a session, on its VM
  agents_api.py --rm SESSION        delete a session and its VM

OpenAI runs the agent. This makes a VM for the session with exe.dev's Python
SDK and starts `codex exec-server` there, which runs the agent's commands in
/workspace.

Settings, from the environment or .env beside this file:
  OPENAI_EXECUTOR_API_KEY  an environment key from the Agents tab of the
                  OpenAI platform; the session VM gets it as CODEX_API_KEY
  OPENAI_BASE_URL default https://openai.int.exe.xyz/v1, an http-proxy
                  integration that adds your OpenAI API key
  OPENAI_API_KEY  use api.openai.com with this key instead
  EXE_API         default https://openai-agents.int.exe.xyz, an exe-dev
                  integration; or EXE_TOKEN, an exe.dev API key
  MODEL           default gpt-6-astra

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

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

import exedev
from openai import AsyncOpenAI

HERE = Path(__file__).resolve().parent
WORKSPACE = "/workspace"
INSTRUCTIONS = """\
You work on an exe.dev VM of your own, as exedev with passwordless sudo, in
/workspace. Run `hostname` for its name. To leave a server running, start it with
`sudo systemd-run --unit=<name> --uid=exedev --working-directory=$PWD <command>`,
on 0.0.0.0 and a port from 3000 to 9999. It is then at
https://<hostname>.exe.xyz:<port>/ for the user, after exe.dev login. Give them
that link."""
FAILED = {"error", "agent.session.environment.failed", "agent.session.failed",
          "agent.session.turn.failed", "agent.session.turn.cancelled"}


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 exe_client():
    if os.environ.get("EXE_TOKEN"):
        return exedev.AsyncClient(os.environ["EXE_TOKEN"])
    api = os.environ.get("EXE_API", "https://openai-agents.int.exe.xyz")
    return exedev.AsyncClient("integration", base_url=api)  # it adds the key


def openai_client():
    if os.environ.get("OPENAI_API_KEY") and not os.environ.get("OPENAI_BASE_URL"):
        return AsyncOpenAI(timeout=360)
    base = os.environ.get("OPENAI_BASE_URL", "https://openai.int.exe.xyz/v1")
    return AsyncOpenAI(base_url=base, api_key="integration", timeout=360)


async def find_vm(exe, session_id):
    for vm in await exe.vms():
        if vm.comment == f"agents-api {session_id}":
            return vm
    return None


async def start_executor(vm, session, key):
    """Starts codex exec-server on the VM, under systemd, for this session."""
    if (await vm.run("systemctl is-active --quiet codex-executor", timeout=30)).ok:
        return
    await vm.run("sudo systemctl reset-failed codex-executor 2>/dev/null; true", timeout=30)
    await vm.run(f"sudo install -d -o exedev -g exedev {WORKSPACE}", check=True, timeout=60)
    await vm.write_file("/home/exedev/.codex-executor.env", f"CODEX_API_KEY={key}\n".encode(), mode=0o600)
    cmd = ["sudo", "systemd-run", "--unit=codex-executor", "--uid=exedev",
           f"--working-directory={WORKSPACE}",
           "--property=EnvironmentFile=/home/exedev/.codex-executor.env",
           "codex", "exec-server", "--remote", session.environment.remote_url,
           "--environment-id", session.environment.id]
    await vm.run(shlex.join(cmd), check=True, timeout=60)


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("--rm", metavar="SESSION", help="delete a session and its VM")
    a = p.parse_args()
    load_env()
    async with openai_client() as client, exe_client() as exe:
        if a.rm:
            vm = await find_vm(exe, a.rm)
            if vm:
                await vm.delete()
                print(f"deleted VM {vm.name}")
            await client.beta.agents.sessions.delete(a.rm)
            print(f"deleted session {a.rm}")
            return
        if not a.prompt:
            p.print_usage()
            return
        key = os.environ.get("OPENAI_EXECUTOR_API_KEY") or sys.exit("agents_api: set OPENAI_EXECUTOR_API_KEY")
        if a.session:
            session = await client.beta.agents.sessions.retrieve(a.session)
            vm = await find_vm(exe, session.id) or sys.exit(f"agents_api: no VM for {session.id}")
        else:
            session = await client.beta.agents.sessions.create(
                agent={"model": os.environ.get("MODEL", "gpt-6-astra"), "instructions": INSTRUCTIONS},
                environment={"type": "self_hosted", "workspace_directory": WORKSPACE},
            )
            vm = await exe.new(comment=f"agents-api {session.id}")
            print(f"session {session.id}, on VM {vm.name}", file=sys.stderr)
        await start_executor(vm, session, key)
        async with client.beta.agents.sessions.stream(session.id, input=a.prompt) as events:
            async for event in events:
                if event.type in FAILED:
                    sys.exit(f"agents_api: {event.type}: {event.to_json(indent=None)}")
                if event.type == "agent.session.turn.output_text.delta":
                    print(event.delta, end="", flush=True)
        print(f"\n\ncontinue: {sys.argv[0]} -s {session.id} \"...\"", file=sys.stderr)


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