OpenEnv

import asyncio
import re
from dataclasses import dataclass

from anthropic import AsyncAnthropic
from anthropic.types import TextBlock
from superserve import AsyncSandbox, NetworkConfig
from textarena_env import TextArenaAction, TextArenaEnv

MODEL = "claude-sonnet-5-5"
SYSTEM = (
    "You are playing Wordle. Reason briefly, then end your reply with one guess: "
    "a five-letter English word in square brackets, for example [crane]."
)

anthropic = AsyncAnthropic()


@dataclass
class Turn:
    guess: str
    reward: float
    signals: dict[str, float]


async def choose_guess(transcript: str) -> str:
    message = await anthropic.messages.create(
        model=MODEL,
        max_tokens=1024,
        system=SYSTEM,
        messages=[{"role": "user", "content": transcript}],
    )
    reply = "".join(b.text for b in message.content if isinstance(b, TextBlock))
    guesses = re.findall(r"\[[a-z]{5}\]", reply.lower())
    return guesses[-1] if guesses else reply


async def run_episode(episode: int) -> list[Turn]:
    sandbox = await AsyncSandbox.create(
        name=f"wordle-episode-{episode}",
        from_template="openenv-wordle",
        preview_access="private",
        network=NetworkConfig(deny_out=["0.0.0.0/0"]),
    )
    try:
        # Start the server in each sandbox so every episode draws its own word.
        async with await sandbox.commands.spawn(
            "uvicorn server.app:app --host 0.0.0.0 --port 8000", cwd="/app/env"
        ):
            await sandbox.commands.run(
                "until curl -fsS http://localhost:8000/health; do sleep 1; done",
                timeout_seconds=60,
            )
            await sandbox.publish_preview_port(8000)
            credential = await sandbox.get_preview_token(8000)
            async with TextArenaEnv(
                base_url=sandbox.get_preview_url(8000),
                headers={credential.header: credential.token},
            ) as env:
                result = await env.reset()
                turns: list[Turn] = []
                while not result.done:
                    # The last message holds the full game transcript so far.
                    transcript = result.observation.messages[-1].content
                    guess = await choose_guess(transcript)
                    result = await env.step(TextArenaAction(message=guess))
                    turn = Turn(
                        guess=guess,
                        reward=result.reward or 0.0,
                        signals=result.observation.info.get("reward_signals", {}),
                    )
                    turns.append(turn)
                    print(f"episode {episode} turn {len(turns)}: {turn}")
                return turns
    finally:
        await sandbox.kill()


async def main() -> None:
    episodes = await asyncio.gather(*(run_episode(i) for i in range(4)))
    for i, turns in enumerate(episodes):
        print(f"episode {i}: {len(turns)} turns, final reward {turns[-1].reward}")


asyncio.run(main())