# main.py - FastAPI application for Pokemon Livestream
import asyncio
import os
import random
import time
import traceback
import logging
from typing import List, Dict, Optional, Set
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
# --- Imports for poke_env and agents ---
from poke_env.player import Player
from poke_env import AccountConfiguration, ServerConfiguration
from poke_env.environment.battle import Battle
# Import the actual agent classes
from agents import OpenAIAgent, GeminiAgent, MistralAgent
# --- Configuration ---
CUSTOM_SERVER_URL = "wss://jofthomas.com/showdown/websocket"
CUSTOM_ACTION_URL = 'https://play.pokemonshowdown.com/action.php?'
CUSTOM_BATTLE_VIEW_URL_TEMPLATE = "https://jofthomas.com/play.pokemonshowdown.com/testclient.html#{battle_id}"
custom_config = ServerConfiguration(CUSTOM_SERVER_URL, CUSTOM_ACTION_URL)
DEFAULT_BATTLE_FORMAT = "gen9randombattle"
# Define available agents with their corresponding classes
AGENT_CONFIGS = {
"OpenAIAgent": {"class": OpenAIAgent, "password_env_var": "OPENAI_AGENT_PASSWORD"},
"GeminiAgent": {"class": GeminiAgent, "password_env_var": "GEMINI_AGENT_PASSWORD"},
"MistralAgent": {"class": MistralAgent, "password_env_var": "MISTRAL_AGENT_PASSWORD"},
}
# Filter out agents with missing passwords
AVAILABLE_AGENT_NAMES = [
name for name, cfg in AGENT_CONFIGS.items()
if os.environ.get(cfg.get("password_env_var", ""))
]
if not AVAILABLE_AGENT_NAMES:
print("FATAL ERROR: No agent configurations have their required password environment variables set. Exiting.")
exit(1)
# --- Global State Variables ---
active_agent_name: Optional[str] = None
active_agent_instance: Optional[Player] = None
active_agent_task: Optional[asyncio.Task] = None
current_battle_instance: Optional[Battle] = None
background_task_handle: Optional[asyncio.Task] = None
# --- Create FastAPI app ---
app = FastAPI(title="Pokemon Battle Livestream")
# --- Helper Functions ---
def get_active_battle(agent: Player) -> Optional[Battle]:
"""Returns the first non-finished battle for an agent."""
if agent and agent._battles:
active_battles = [b for b in agent._battles.values() if not b.finished]
if active_battles:
if active_battles[0].battle_tag:
return active_battles[0]
else:
print(f"WARN: Found active battle for {agent.username} but it has no battle_tag yet.")
return None
return None
def create_battle_iframe(battle_id: str) -> str:
"""Creates the HTML for the battle iframe."""
print("CURRENT BATTLE ID: ", battle_id)
battle_url = f"https://jofthomas.com/play.pokemonshowdown.com/testclient.html#{battle_id}"
return f"""
Pokemon Battle Livestream
"""
def create_idle_html(status_message: str, instruction: str) -> str:
"""Creates a visually appealing idle screen HTML."""
return f"""
Pokemon Battle Livestream
{status_message}
{instruction}
"""
def create_error_html(error_msg: str) -> str:
"""Creates HTML to display an error message."""
return f"""
Pokemon Battle Livestream - Error
An Error Occurred
{error_msg}
"""
async def update_display_html(new_html: str) -> None:
"""Updates the current display HTML and broadcasts to all clients."""
await manager.update_all(new_html)
print("HTML Display UPDATED and broadcasted.")
# --- Agent Lifecycle Management ---
async def select_and_activate_new_agent():
"""Selects a random available agent, instantiates it, and starts its listening task."""
global active_agent_name, active_agent_instance, active_agent_task
if not AVAILABLE_AGENT_NAMES:
print("Lifecycle: No available agents with passwords set.")
await update_display_html(create_error_html("No agents available (check password env vars)."))
return False
selected_name = random.choice(AVAILABLE_AGENT_NAMES)
config = AGENT_CONFIGS[selected_name]
AgentClass = config["class"]
password_env_var = config["password_env_var"]
agent_password = os.environ.get(password_env_var)
print(f"Lifecycle: Activating agent '{selected_name}'...")
await update_display_html(create_idle_html("Selecting Next Agent...", f"Preparing {selected_name}..."))
try:
account_config = AccountConfiguration(selected_name, agent_password)
agent = AgentClass(
account_configuration=account_config,
server_configuration=custom_config,
battle_format=DEFAULT_BATTLE_FORMAT,
log_level=logging.INFO,
max_concurrent_battles=1
)
# Start the task to accept exactly one battle challenge
task = asyncio.create_task(agent.accept_challenges(None, 1), name=f"accept_challenge_{selected_name}")
task.add_done_callback(log_task_exception)
# Update global state
active_agent_name = selected_name
active_agent_instance = agent
active_agent_task = task
print(f"Lifecycle: Agent '{selected_name}' is active and listening for 1 challenge.")
await update_display_html(create_idle_html(f"Agent {selected_name} is ready!",
f"Please challenge {selected_name} to a {DEFAULT_BATTLE_FORMAT} battle."))
return True
except Exception as e:
error_msg = f"Failed to activate agent '{selected_name}': {e}"
print(error_msg)
traceback.print_exc()
# Clear state if activation failed
active_agent_name = None
active_agent_instance = None
active_agent_task = None
return False
async def check_for_new_battle():
"""Checks if the active agent has started a battle."""
global active_agent_instance, current_battle_instance
if active_agent_instance:
battle = get_active_battle(active_agent_instance)
if battle and battle.battle_tag:
print(f"Lifecycle: Agent '{active_agent_name}' started battle: {battle.battle_tag}")
current_battle_instance = battle
# Prevent accepting more challenges
if active_agent_task and not active_agent_task.done():
print(f"Lifecycle: Cancelling accept_challenges task for {active_agent_name} as battle started.")
active_agent_task.cancel()
async def deactivate_current_agent(reason: str = "cycle"):
"""Cleans up the currently active agent and resets state."""
global active_agent_name, active_agent_instance, active_agent_task, current_battle_instance
print(f"Lifecycle: Deactivating agent '{active_agent_name}' (Reason: {reason})...")
# Different messages based on deactivation reason
if reason == "battle_end":
await update_display_html(create_idle_html("Battle Finished", "Preparing to select next agent..."))
elif reason == "cycle":
await update_display_html(create_idle_html("Cycling Agents", "Preparing for next selection..."))
else:
await update_display_html(create_idle_html(f"Agent Reset ({reason})", "Preparing for next selection..."))
agent = active_agent_instance
task = active_agent_task
# Store a local copy of the battle instance before clearing it
last_battle_instance = current_battle_instance
# Clear state first to prevent race conditions
active_agent_name = None
active_agent_instance = None
active_agent_task = None
current_battle_instance = None
# Cancel the accept_challenges task if it's still running
if task and not task.done():
print(f"Lifecycle: Cancelling task for {agent.username if agent else 'unknown agent'}...")
task.cancel()
try:
await asyncio.wait_for(task, timeout=2.0)
except asyncio.CancelledError:
print(f"Lifecycle: Task cancellation confirmed.")
except asyncio.TimeoutError:
print(f"Lifecycle: Task did not confirm cancellation within timeout.")
except Exception as e:
print(f"Lifecycle: Error during task cancellation wait: {e}")
# Disconnect the player
if agent:
print(f"Lifecycle: Disconnecting player {agent.username}...")
try:
if hasattr(agent, '_websocket') and agent._websocket and agent._websocket.open:
await agent.disconnect()
print(f"Lifecycle: Player {agent.username} disconnected.")
else:
print(f"Lifecycle: Player {agent.username} already disconnected or websocket not ready.")
except Exception as e:
print(f"ERROR during agent disconnect ({agent.username}): {e}")
# Add a brief delay after deactivation to ensure clean state
await asyncio.sleep(3)
print(f"Lifecycle: Agent deactivation complete.")
# --- Main Background Task ---
async def manage_agent_lifecycle():
"""Runs the main loop selecting, running, and cleaning up agents sequentially."""
global active_agent_instance, active_agent_task, current_battle_instance
print("Background lifecycle manager started.")
REFRESH_INTERVAL_SECONDS = 3
loop_counter = 0
while True:
loop_counter += 1
try:
print(f"\n--- Lifecycle Check #{loop_counter} [{time.strftime('%H:%M:%S')}] ---")
# ==================================
# State 1: No agent active
# ==================================
if active_agent_instance is None:
print(f"[{loop_counter}] State 1: No active agent. Selecting...")
activated = await select_and_activate_new_agent()
if not activated:
print(f"[{loop_counter}] State 1: Activation failed. Waiting before retry.")
await asyncio.sleep(10)
continue
else:
print(f"[{loop_counter}] State 1: Agent '{active_agent_name}' activated.")
# ==================================
# State 2: Agent is active
# ==================================
else:
print(f"[{loop_counter}] State 2: Agent '{active_agent_name}' active.")
# --- Sub-state: Monitoring an ongoing battle ---
if current_battle_instance is not None:
battle_tag = current_battle_instance.battle_tag
print(f"[{loop_counter}] State 2: Monitoring battle {battle_tag}")
# Get the potentially updated battle object from the agent's internal list
battle_obj = active_agent_instance._battles.get(battle_tag)
if battle_obj and battle_obj.finished:
# Battle is finished, deactivate agent
print(f"[{loop_counter}] Battle {battle_tag} is FINISHED. Deactivating agent.")
await deactivate_current_agent(reason="battle_end")
await asyncio.sleep(5) # Brief cooldown between battles
continue
elif not battle_obj:
# Battle object missing unexpectedly, deactivate agent
print(f"[{loop_counter}] WARNING: Battle object for {battle_tag} MISSING from agent's list! Deactivating.")
await deactivate_current_agent(reason="battle_object_missing")
continue
else:
# Battle is ongoing and battle object exists
pass
# --- Sub-state: Checking for a new battle / Agent is idle ---
elif current_battle_instance is None:
print(f"[{loop_counter}] State 2: Checking for new battle...")
await check_for_new_battle()
# --- If a NEW battle was just detected by check_for_new_battle() ---
if current_battle_instance:
battle_tag = current_battle_instance.battle_tag
print(f"[{loop_counter}] State 2: *** NEW BATTLE DETECTED: {battle_tag} ***")
# Check if the battle ID has the non-public/suffixed format
parts = battle_tag.split('-')
is_suffixed_format = len(parts) > 3 and parts[2].isdigit()
if is_suffixed_format:
# Non-public format detected, forfeit immediately
print(f"[{loop_counter}] Detected non-public battle format ({battle_tag}). Forfeiting.")
await update_display_html(create_idle_html(f"Forfeiting non-public battle...", f"Agent {active_agent_name} detected private room."))
try:
# Send forfeit command TO THE SPECIFIC BATTLE ROOM
await active_agent_instance.send_message("/forfeit", battle_tag)
print(f"[{loop_counter}] Sent /forfeit command for {battle_tag}.")
await asyncio.sleep(1) # Give a moment for forfeit to process server-side
except Exception as forfeit_err:
print(f"[{loop_counter}] ERROR sending forfeit command for {battle_tag}: {forfeit_err}")
# Deactivate the agent as the battle is effectively over
await deactivate_current_agent(reason="forfeited_private_battle")
continue
else:
# Public battle format detected, display the battle iframe
print(f"[{loop_counter}] Public battle format detected. Displaying battle.")
await update_display_html(create_battle_iframe(battle_tag))
# --- If no new battle was found, agent remains idle ---
else:
print(f"[{loop_counter}] State 2: No new battle found. Agent '{active_agent_name}' remains idle.")
expected_idle_html = create_idle_html(f"Agent {active_agent_name} is ready!",
f"Please challenge {active_agent_name} to a {DEFAULT_BATTLE_FORMAT} battle.")
await update_display_html(expected_idle_html)
# --- Global Exception Handling for the main loop ---
except Exception as e:
print(f"ERROR in main lifecycle loop #{loop_counter}: {e}")
traceback.print_exc()
if active_agent_instance:
# If an error happened while an agent was active, try to deactivate it
await deactivate_current_agent(reason="main_loop_error")
else:
# If error happened before agent activation, show generic error
await update_display_html(create_error_html(f"Error in lifecycle manager: {e}"))
# Wait longer after a major error before trying again
await asyncio.sleep(10)
# --- Sleep at the end of each loop iteration (if no 'continue' was hit) ---
print(f"[{loop_counter}] Sleeping {REFRESH_INTERVAL_SECONDS}s before next check.")
await asyncio.sleep(REFRESH_INTERVAL_SECONDS)
def log_task_exception(task: asyncio.Task):
"""Callback to log exceptions from background tasks."""
try:
if task.cancelled():
print(f"Task {task.get_name()} was cancelled.")
return
task.result() # Raises exception if task failed
print(f"Task {task.get_name()} finished cleanly.")
except asyncio.CancelledError:
pass
except Exception as e:
print(f"Exception in background task {task.get_name()}: {e}")
traceback.print_exc()
# --- WebSocket connection manager ---
class ConnectionManager:
def __init__(self):
self.active_connections: Set[WebSocket] = set()
self.current_html: str = create_idle_html("Initializing...", "Setting up Pokemon Battle Stream")
async def connect(self, websocket: WebSocket):
await websocket.accept()
self.active_connections.add(websocket)
# Send current state to newly connected client
await websocket.send_text(self.current_html)
def disconnect(self, websocket: WebSocket):
self.active_connections.remove(websocket)
async def update_all(self, html: str):
"""Update the current HTML and broadcast to all clients"""
self.current_html = html
if self.active_connections:
# Only log if there are connections to update
print(f"Broadcasting update to {len(self.active_connections)} clients")
# Make a copy of the connections set to avoid modification during iteration
connections_copy = self.active_connections.copy()
for connection in connections_copy:
try:
await connection.send_text(html)
except Exception as e:
print(f"Error sending to client: {e}")
# Don't remove here - will be handled by disconnect route
manager = ConnectionManager()
# --- API Routes ---
@app.get("/", response_class=HTMLResponse)
async def get_homepage():
"""Serves the main HTML page with WebSocket connection."""
return """
Pokemon Battle Livestream
"""
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
try:
while True:
# Just keep the connection open
# The updates are pushed from the lifecycle manager
await websocket.receive_text()
except WebSocketDisconnect:
manager.disconnect(websocket)
except Exception as e:
print(f"WebSocket error: {e}")
manager.disconnect(websocket)
@app.on_event("startup")
async def startup_event():
"""Start background tasks when the application starts."""
global background_task_handle
# Mount static files directory for images
app.mount("/static", StaticFiles(directory="static"), name="static")
print("🚀 Starting background tasks")
background_task_handle = asyncio.create_task(manage_agent_lifecycle(), name="lifecycle_manager")
background_task_handle.add_done_callback(log_task_exception)
print("✅ Background tasks started")
@app.on_event("shutdown")
async def shutdown_event():
"""Clean up tasks when shutting down."""
global background_task_handle, active_agent_instance
print("\nShutting down app. Cleaning up...")
# Cancel background task
if background_task_handle and not background_task_handle.done():
background_task_handle.cancel()
try:
await asyncio.wait_for(background_task_handle, timeout=5.0)
except (asyncio.CancelledError, asyncio.TimeoutError):
pass
# Disconnect agent if active
if active_agent_instance:
try:
if hasattr(active_agent_instance, '_websocket') and active_agent_instance._websocket and active_agent_instance._websocket.open:
await active_agent_instance.disconnect()
except Exception as e:
print(f"Error during agent disconnect on shutdown: {e}")
print("Cleanup complete. Shutting down.")
# For direct script execution
if __name__ == "__main__":
import uvicorn
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.getLogger('poke_env').setLevel(logging.WARNING)
print("Starting Pokemon Battle Livestream Server...")
print("="*50)
if not AVAILABLE_AGENT_NAMES:
print("FATAL: No agents found with configured passwords. Please set environment variables:")
for name, cfg in AGENT_CONFIGS.items():
print(f"- {cfg.get('password_env_var', 'N/A')} (for agent: {name})")
print("="*50)
exit("Exiting due to missing agent passwords.")
else:
print("Found available agents:")
for name in AVAILABLE_AGENT_NAMES:
print(f"- {name}")
print("="*50)
# Run with uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)