Spaces:

Jofthomas
/

twitch_streaming

Running

App Files Files Community

Jofthomas commited on Apr 16

Commit

4cbc162

verified ·

1 Parent(s): 42e69ad

Update main.py

Browse files

Files changed (1) hide show

main.py +563 -359

main.py CHANGED Viewed

@@ -54,181 +54,73 @@ 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"""
-    <html>
-    <head>
-        <title>Pokemon Battle Livestream</title>
-        <style>
-            body, html {{
-                margin: 0;
-                padding: 0;
-                height: 100%;
-                overflow: hidden;
-            }}
-            iframe {{
-                position: fixed;
-                top: 0;
-                left: 0;
-                bottom: 0;
-                right: 0;
-                width: 100%;
-                height: 100%;
-                border: none;
-                margin: 0;
-                padding: 0;
-                overflow: hidden;
-                z-index: 999999;
-            }}
-        </style>
-    </head>
-    <body>
-        <iframe
-            src="{battle_url}"
-            allowfullscreen
-        ></iframe>
-    </body>
-    </html>
     """
 def create_idle_html(status_message: str, instruction: str) -> str:
-    """Creates a visually appealing idle screen HTML."""
     return f"""
-    <html>
-    <head>
-        <title>Pokemon Battle Livestream</title>
-        <style>
-            body, html {{
-                margin: 0;
-                padding: 0;
-                height: 100%;
-                overflow: hidden;
-                font-family: Arial, sans-serif;
-            }}
-            .container {{
-                position: fixed;
-                top: 0;
-                left: 0;
-                width: 100%;
-                height: 100%;
-                display: flex;
-                flex-direction: column;
-                justify-content: center;
-                align-items: center;
-                background-image: url('/static/pokemon_huggingface.png');
-                background-size: cover;
-                background-position: center;
-                z-index: 999999;
-            }}
-            .message-box {{
-                background-color: rgba(0, 0, 0, 0.65);
-                padding: 30px;
-                border-radius: 15px;
-                max-width: 80%;
-                text-align: center;
-            }}
-            .status {{
-                font-size: 2.5em;
-                margin-bottom: 20px;
-                text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.8);
-                color: white;
-            }}
-            .instruction {{
-                font-size: 1.5em;
-                text-shadow: 1px 1px 3px rgba(0, 0, 0, 0.8);
-                color: white;
-            }}
-        </style>
-    </head>
-    <body>
-        <div class="container">
-            <div class="message-box">
-                <p class="status">{status_message}</p>
-                <p class="instruction">{instruction}</p>
-            </div>
         </div>
-    </body>
-    </html>
     """
 def create_error_html(error_msg: str) -> str:
-    """Creates HTML to display an error message."""
     return f"""
-    <html>
-    <head>
-        <title>Pokemon Battle Livestream - Error</title>
-        <style>
-            body, html {{
-                margin: 0;
-                padding: 0;
-                height: 100%;
-                overflow: hidden;
-                font-family: Arial, sans-serif;
-            }}
-            .container {{
-                position: fixed;
-                top: 0;
-                left: 0;
-                width: 100%;
-                height: 100%;
-                display: flex;
-                flex-direction: column;
-                justify-content: center;
-                align-items: center;
-                background-color: #330000;
-                z-index: 999999;
-            }}
-            .message-box {{
-                background-color: rgba(200, 0, 0, 0.7);
-                padding: 30px;
-                border-radius: 15px;
-                max-width: 80%;
-                text-align: center;
-            }}
-            .title {{
-                font-size: 2em;
-                margin-bottom: 20px;
-                color: white;
-            }}
-            .message {{
-                font-size: 1.2em;
-                color: #ffdddd;
-            }}
-        </style>
-    </head>
-    <body>
-        <div class="container">
-            <div class="message-box">
-                <p class="title">An Error Occurred</p>
-                <p class="message">{error_msg}</p>
-            </div>
         </div>
-    </body>
-    </html>
     """
-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():
@@ -237,7 +129,7 @@ async def select_and_activate_new_agent():
     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)
@@ -247,7 +139,8 @@ async def select_and_activate_new_agent():
     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)
@@ -260,16 +153,18 @@ async def select_and_activate_new_agent():
         )
         # 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 <strong>{selected_name}</strong> is ready!",
                                       f"Please challenge <strong>{selected_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle."))
         return True
@@ -277,7 +172,8 @@ async def select_and_activate_new_agent():
         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
@@ -285,74 +181,103 @@ async def select_and_activate_new_agent():
         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():
@@ -360,14 +285,19 @@ async def manage_agent_lifecycle():
     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
             # ==================================
@@ -375,214 +305,425 @@ async def manage_agent_lifecycle():
                 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 <strong>{active_agent_name}</strong> is ready!",
-                                                             f"Please challenge <strong>{active_agent_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> 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 """
     <!DOCTYPE html>
-    <html>
     <head>
         <title>Pokemon Battle Livestream</title>
         <style>
-            body, html {
                 margin: 0;
                 padding: 0;
                 height: 100%;
-                overflow: hidden;
             }
             #stream-container {
-                position: fixed;
                 top: 0;
                 left: 0;
-                bottom: 0;
-                right: 0;
                 width: 100%;
                 height: 100%;
             }
         </style>
     </head>
     <body>
-        <div id="stream-container"></div>
         <script>
             const streamContainer = document.getElementById('stream-container');
-            // Create WebSocket connection
-            const ws = new WebSocket(`${location.protocol === 'https:' ? 'wss' : 'ws'}://${location.host}/ws`);
-            // Connection opened
-            ws.addEventListener('open', (event) => {
-                console.log('Connected to WebSocket server');
-            });
-            // Listen for messages
-            ws.addEventListener('message', (event) => {
-                console.log('Received update from server');
-                // Set inner HTML directly - complete document replacement
-                const parser = new DOMParser();
-                const doc = parser.parseFromString(event.data, 'text/html');
-                const bodyContent = doc.body.innerHTML;
-                // Update just the body content to avoid page refresh flicker
-                if (bodyContent) {
-                    streamContainer.innerHTML = bodyContent;
-                }
-            });
-            // Handle connection close/errors
-            ws.addEventListener('close', (event) => {
-                console.log('WebSocket connection closed');
-                setTimeout(() => {
-                    console.log('Attempting to reconnect...');
-                    location.reload();
-                }, 3000);
-            });
-            ws.addEventListener('error', (event) => {
-                console.error('WebSocket error:', event);
-            });
         </script>
     </body>
     </html>
@@ -593,25 +734,45 @@ 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")
@@ -619,48 +780,91 @@ async def startup_event():
 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=7860)

 # --- 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:
+            # Ensure the battle object has a battle_tag before returning
+            if hasattr(active_battles[0], 'battle_tag') and active_battles[0].battle_tag:
+                 # Check if the battle_tag has the expected format (starts with 'battle-')
+                if active_battles[0].battle_tag.startswith("battle-"):
+                    return active_battles[0]
+                else:
+                    # This handles cases where the battle object might exist but tag isn't ready
+                    # print(f"DEBUG: Found active battle for {agent.username} but tag '{active_battles[0].battle_tag}' not ready.")
+                    return None
             else:
+                # print(f"DEBUG: Found active battle for {agent.username} but it has no battle_tag attribute yet.")
                 return None
     return None
 def create_battle_iframe(battle_id: str) -> str:
+    """Creates JUST the HTML for the battle iframe tag."""
+    print("Creating iframe content for battle ID: ", battle_id)
+    # Use the official client URL unless you specifically need the test client
+    # battle_url = f"https://play.pokemonshowdown.com/{battle_id}"
+    battle_url = f"https://jofthomas.com/play.pokemonshowdown.com/testclient.html#{battle_id}" # Using your custom URL
+    # Return ONLY the iframe tag with a class for styling
     return f"""
+    <iframe
+        id="battle-iframe"
+        class="battle-iframe"
+        src="{battle_url}"
+        allowfullscreen
+    ></iframe>
     """
 def create_idle_html(status_message: str, instruction: str) -> str:
+    """Creates a visually appealing idle screen HTML fragment."""
+    # Returns ONLY the content div, not the full HTML page
     return f"""
+    <div class="content-container idle-container">
+        <div class="message-box">
+            <p class="status">{status_message}</p>
+            <p class="instruction">{instruction}</p>
         </div>
+    </div>
     """
 def create_error_html(error_msg: str) -> str:
+    """Creates HTML fragment to display an error message."""
+     # Returns ONLY the content div, not the full HTML page
     return f"""
+    <div class="content-container error-container">
+        <div class="message-box">
+            <p class="status">🚨 Error 🚨</p>
+            <p class="instruction">{error_msg}</p>
         </div>
+    </div>
     """
+async def update_display_html(new_html_fragment: str) -> None:
+    """Updates the current display HTML fragment and broadcasts to all clients."""
+    # Pass the fragment directly
+    await manager.update_all(new_html_fragment)
+    print("HTML Display FRAGMENT UPDATED and broadcasted.")
 # --- Agent Lifecycle Management ---
 async def select_and_activate_new_agent():
     if not AVAILABLE_AGENT_NAMES:
         print("Lifecycle: No available agents with passwords set.")
+        await update_display_html(create_error_html("No agents available. Check server logs/environment variables."))
         return False
     selected_name = random.choice(AVAILABLE_AGENT_NAMES)
     agent_password = os.environ.get(password_env_var)
     print(f"Lifecycle: Activating agent '{selected_name}'...")
+    # Use HTML tags for slight emphasis if desired
+    await update_display_html(create_idle_html("Selecting Next Agent...", f"Preparing <strong>{selected_name}</strong>..."))
     try:
         account_config = AccountConfiguration(selected_name, agent_password)
         )
         # Start the task to accept exactly one battle challenge
+        # Setting name for easier debugging
+        task = asyncio.create_task(agent.accept_challenges(None, 1), name=f"AcceptChallenge_{selected_name}")
+        task.add_done_callback(log_task_exception) # Add callback for errors
         # 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.")
+        # Use HTML tags for slight emphasis
+        await update_display_html(create_idle_html(f"Agent Ready: <strong>{selected_name}</strong>",
                                       f"Please challenge <strong>{selected_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle."))
         return True
         error_msg = f"Failed to activate agent '{selected_name}': {e}"
         print(error_msg)
         traceback.print_exc()
+        await update_display_html(create_error_html(f"Error activating {selected_name}. Please wait or check logs."))
         # Clear state if activation failed
         active_agent_name = None
         active_agent_instance = None
         return False
 async def check_for_new_battle():
+    """Checks if the active agent has started a battle with a valid tag."""
     global active_agent_instance, current_battle_instance
     if active_agent_instance:
         battle = get_active_battle(active_agent_instance)
+        # Check if battle exists AND has a valid battle_tag
         if battle and battle.battle_tag:
             print(f"Lifecycle: Agent '{active_agent_name}' started battle: {battle.battle_tag}")
             current_battle_instance = battle
+            # Stop the agent from listening for more challenges once a battle starts
             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()
+                # Optional: Wait briefly for cancellation confirmation, but don't block excessively
+                # try:
+                #     await asyncio.wait_for(active_agent_task, timeout=0.5)
+                # except (asyncio.CancelledError, asyncio.TimeoutError):
+                #     pass # Expected outcomes
+        # else:
+            # print(f"DEBUG: get_active_battle returned None or battle without tag.")
 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
+    agent_name_to_deactivate = active_agent_name # Store before clearing
+    print(f"Lifecycle: Deactivating agent '{agent_name_to_deactivate}' (Reason: {reason})...")
+    # Display appropriate intermediate message
     if reason == "battle_end":
+        await update_display_html(create_idle_html("Battle Finished!", f"Agent <strong>{agent_name_to_deactivate}</strong> completed the match."))
     elif reason == "cycle":
+        await update_display_html(create_idle_html("Cycling Agents", f"Switching from <strong>{agent_name_to_deactivate}</strong>..."))
+    elif reason == "forfeited_private_battle":
+         await update_display_html(create_idle_html("Switching Agent", f"Agent <strong>{agent_name_to_deactivate}</strong> forfeited a private battle."))
+    else: # Generic reason or error
+        await update_display_html(create_idle_html(f"Resetting Agent ({reason})", f"Cleaning up <strong>{agent_name_to_deactivate}</strong>..."))
+    # Give users a moment to see the intermediate message
+    await asyncio.sleep(3) # Adjust duration as needed
+    # Show the "preparing next agent" message before lengthy cleanup
+    await update_display_html(create_idle_html("Preparing Next Agent...", "Please wait..."))
     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 # Not strictly needed now
+    # --- Crucial: Clear global state variables FIRST ---
+    # This prevents race conditions where the lifecycle loop might try to
+    # access the agent while it's being deactivated.
     active_agent_name = None
     active_agent_instance = None
     active_agent_task = None
     current_battle_instance = None
+    print(f"Lifecycle: Global state cleared for '{agent_name_to_deactivate}'.")
+    # --- Now perform cleanup actions ---
+    # Cancel the accept_challenges task if it's still running (it might already be done/cancelled)
     if task and not task.done():
+        print(f"Lifecycle: Ensuring task cancellation for {agent_name_to_deactivate} ({task.get_name()})...")
         task.cancel()
         try:
+            # Wait briefly for the task to acknowledge cancellation
             await asyncio.wait_for(task, timeout=2.0)
+            print(f"Lifecycle: Task cancellation confirmed for {agent_name_to_deactivate}.")
         except asyncio.CancelledError:
+             print(f"Lifecycle: Task cancellation confirmation (CancelledError) for {agent_name_to_deactivate}.")
         except asyncio.TimeoutError:
+            print(f"Lifecycle: Task did not confirm cancellation within timeout for {agent_name_to_deactivate}.")
         except Exception as e:
+            # Catch other potential errors during task cleanup
+            print(f"Lifecycle: Error during task cancellation wait for {agent_name_to_deactivate}: {e}")
+    # Disconnect the player (ensure agent object exists)
     if agent:
         print(f"Lifecycle: Disconnecting player {agent.username}...")
         try:
+            # Check websocket state before attempting disconnection
             if hasattr(agent, '_websocket') and agent._websocket and agent._websocket.open:
+                 await agent.disconnect()
+                 print(f"Lifecycle: Player {agent.username} disconnected successfully.")
             else:
+                 print(f"Lifecycle: Player {agent.username} already disconnected or websocket not available.")
         except Exception as e:
+            # Log errors during disconnection but don't halt the process
             print(f"ERROR during agent disconnect ({agent.username}): {e}")
+            traceback.print_exc() # Log full traceback for debugging
+    # Add a brief delay AFTER deactivation before the loop potentially selects a new agent
+    await asyncio.sleep(2) # Reduced from 3, adjust as needed
+    print(f"Lifecycle: Agent '{agent_name_to_deactivate}' deactivation complete.")
 # --- Main Background Task ---
 async def manage_agent_lifecycle():
     global active_agent_instance, active_agent_task, current_battle_instance
     print("Background lifecycle manager started.")
+    REFRESH_INTERVAL_SECONDS = 3 # How often to check state when idle/in battle
+    LOOP_COOLDOWN_SECONDS = 1 # Small delay at end of loop if no other waits occurred
+    ERROR_RETRY_DELAY_SECONDS = 10 # Longer delay after errors
+    POST_BATTLE_DELAY_SECONDS = 5 # Delay after a battle finishes before selecting next agent
     loop_counter = 0
     while True:
         loop_counter += 1
+        loop_start_time = time.monotonic()
+        print(f"\n--- Lifecycle Check #{loop_counter} [{time.strftime('%H:%M:%S')}] ---")
+        try:
             # ==================================
             # State 1: No agent active
             # ==================================
                 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 {ERROR_RETRY_DELAY_SECONDS}s before retry.")
+                    await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
                 else:
+                    print(f"[{loop_counter}] State 1: Agent '{active_agent_name}' activated successfully.")
+                    # No sleep here, proceed to next check immediately if needed
             # ==================================
             # State 2: Agent is active
             # ==================================
             else:
+                agent_name = active_agent_name # Cache for logging
+                print(f"[{loop_counter}] State 2: Agent '{agent_name}' is active.")
+                # --- Sub-state: Check for new battle if none is tracked ---
+                if current_battle_instance is None:
+                    print(f"[{loop_counter}] State 2a: Checking for new battle for '{agent_name}'...")
+                    await check_for_new_battle() # This updates current_battle_instance if found
                     if current_battle_instance:
                         battle_tag = current_battle_instance.battle_tag
+                        print(f"[{loop_counter}] State 2a: *** NEW BATTLE DETECTED: {battle_tag} for '{agent_name}' ***")
+                        # Check for non-public/suffixed format (heuristic: more than 3 parts, 3rd part is number)
                         parts = battle_tag.split('-')
                         is_suffixed_format = len(parts) > 3 and parts[2].isdigit()
                         if is_suffixed_format:
+                            # Forfeit immediately if it looks like a private/suffixed battle ID
+                            print(f"[{loop_counter}] Detected potentially non-public battle format ({battle_tag}). Forfeiting.")
+                            # Don't update display yet, do it before deactivation
                             try:
+                                if active_agent_instance: # Ensure agent still exists
+                                    await active_agent_instance.forfeit(battle_tag)
+                                    # await active_agent_instance.send_message("/forfeit", battle_tag) # Alternative
+                                    print(f"[{loop_counter}] Sent forfeit command for {battle_tag}.")
+                                    await asyncio.sleep(1.5) # Give forfeit time to register
                             except Exception as forfeit_err:
+                                print(f"[{loop_counter}] ERROR sending forfeit for {battle_tag}: {forfeit_err}")
+                            # Deactivate agent after forfeit attempt
                             await deactivate_current_agent(reason="forfeited_private_battle")
+                            continue # Skip rest of the loop for this iteration
                         else:
+                            # Public battle format - display the iframe
+                            print(f"[{loop_counter}] Public battle format detected. Displaying battle {battle_tag}.")
                             await update_display_html(create_battle_iframe(battle_tag))
+                            # Now fall through to monitor this battle in the next section
                     else:
+                        # No new battle found, agent remains idle
+                        print(f"[{loop_counter}] State 2a: No new battle found. Agent '{agent_name}' remains idle, waiting for challenge.")
+                        # Periodically refresh idle screen to ensure consistency
+                        idle_html = create_idle_html(f"Agent Ready: <strong>{agent_name}</strong>",
+                                                     f"Please challenge <strong>{agent_name}</strong> to a <strong>{DEFAULT_BATTLE_FORMAT}</strong> battle.")
+                        await update_display_html(idle_html)
+                        await asyncio.sleep(REFRESH_INTERVAL_SECONDS) # Wait before next check if idle
+                # --- Sub-state: Monitor ongoing battle ---
+                if current_battle_instance is not None:
+                    battle_tag = current_battle_instance.battle_tag
+                    print(f"[{loop_counter}] State 2b: Monitoring battle {battle_tag} for '{agent_name}'")
+                    # Ensure agent instance still exists before accessing its battles
+                    if not active_agent_instance:
+                        print(f"[{loop_counter}] WARNING: Agent instance for '{agent_name}' disappeared while monitoring battle {battle_tag}! Deactivating.")
+                        await deactivate_current_agent(reason="agent_disappeared_mid_battle")
+                        continue
+                    # Get potentially updated battle object directly from agent's state
+                    # Use .get() for safety
+                    battle_obj = active_agent_instance._battles.get(battle_tag)
+                    if battle_obj and battle_obj.finished:
+                        print(f"[{loop_counter}] Battle {battle_tag} is FINISHED. Deactivating agent '{agent_name}'.")
+                        await deactivate_current_agent(reason="battle_end")
+                        print(f"[{loop_counter}] Waiting {POST_BATTLE_DELAY_SECONDS}s post-battle before selecting next agent.")
+                        await asyncio.sleep(POST_BATTLE_DELAY_SECONDS)
+                        continue # Start next loop iteration to select new agent
+                    elif not battle_obj:
+                        # This can happen briefly during transitions or if battle ends unexpectedly
+                        print(f"[{loop_counter}] WARNING: Battle object for {battle_tag} not found in agent's list for '{agent_name}'. Battle might have ended abruptly. Deactivating.")
+                        await deactivate_current_agent(reason="battle_object_missing")
+                        continue
+                    else:
+                        # Battle is ongoing, battle object exists, iframe should be displayed
+                        print(f"[{loop_counter}] Battle {battle_tag} ongoing for '{agent_name}'.")
+                        # Optionally: Could re-send iframe HTML periodically if needed, but usually not necessary
+                        # await update_display_html(create_battle_iframe(battle_tag))
+                        await asyncio.sleep(REFRESH_INTERVAL_SECONDS) # Wait before next check
         # --- Global Exception Handling for the main loop ---
+        except asyncio.CancelledError:
+            print("Lifecycle manager task cancelled.")
+            raise # Re-raise to ensure proper shutdown
         except Exception as e:
+            print(f"!!! ERROR in main lifecycle loop #{loop_counter}: {e} !!!")
             traceback.print_exc()
+            current_agent_name = active_agent_name # Cache name before deactivation attempts
             if active_agent_instance:
+                print(f"Attempting to deactivate agent '{current_agent_name}' due to loop error...")
+                try:
+                    await deactivate_current_agent(reason="main_loop_error")
+                except Exception as deactivation_err:
+                    print(f"Error during error-handling deactivation: {deactivation_err}")
+                    # Ensure state is cleared even if deactivation fails partially
+                    active_agent_name = None
+                    active_agent_instance = None
+                    active_agent_task = None
+                    current_battle_instance = None
             else:
+                # Error happened potentially before agent activation or after clean deactivation
+                print("No active agent instance during loop error.")
+                # Show a generic error on the frontend
+                await update_display_html(create_error_html(f"A server error occurred in the lifecycle manager. Please wait. ({e})"))
             # Wait longer after a major error before trying again
+            print(f"Waiting {ERROR_RETRY_DELAY_SECONDS}s after loop error.")
+            await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
+            continue # Go to next loop iteration after error handling
+        # --- Delay at end of loop if no other significant waits happened ---
+        elapsed_time = time.monotonic() - loop_start_time
+        if elapsed_time < LOOP_COOLDOWN_SECONDS:
+             await asyncio.sleep(LOOP_COOLDOWN_SECONDS - elapsed_time)
 def log_task_exception(task: asyncio.Task):
+    """Callback to log exceptions from background tasks (like accept_challenges)."""
     try:
         if task.cancelled():
+            # Don't log cancellation as an error, it's often expected
+            print(f"Task '{task.get_name()}' was cancelled.")
             return
+        # Accessing result will raise exception if task failed
+        task.result()
+        print(f"Task '{task.get_name()}' completed successfully.")
     except asyncio.CancelledError:
+         print(f"Task '{task.get_name()}' confirmed cancelled (exception caught).")
+         pass # Expected
     except Exception as e:
+        # Log actual errors
+        print(f"!!! Exception in background task '{task.get_name()}': {e} !!!")
         traceback.print_exc()
+        # Optionally: Trigger some recovery or notification here if needed
 # --- WebSocket connection manager ---
 class ConnectionManager:
     def __init__(self):
         self.active_connections: Set[WebSocket] = set()
+        # Initialize with the idle HTML fragment
+        self.current_html_fragment: str = create_idle_html("Initializing...", "Setting up Pokémon Battle Stream")
     async def connect(self, websocket: WebSocket):
         await websocket.accept()
         self.active_connections.add(websocket)
+        print(f"Client connected. Sending current state. Total clients: {len(self.active_connections)}")
+        # Send current state (HTML fragment) to newly connected client
+        try:
+            await websocket.send_text(self.current_html_fragment)
+        except Exception as e:
+             print(f"Error sending initial state to new client: {e}")
+             # Consider removing the connection if initial send fails
+             await self.disconnect(websocket)
+    async def disconnect(self, websocket: WebSocket):
+        # Use discard() to safely remove even if not present
+        self.active_connections.discard(websocket)
+        print(f"Client disconnected. Total clients: {len(self.active_connections)}")
+    async def update_all(self, html_fragment: str):
+        """Update the current HTML fragment and broadcast to all clients."""
+        if self.current_html_fragment == html_fragment:
+             # print("Skipping broadcast, HTML fragment unchanged.")
+             return # Avoid unnecessary updates if content is identical
+        self.current_html_fragment = html_fragment
+        if not self.active_connections:
+             # print("No active connections to broadcast update to.")
+             return
+        print(f"Broadcasting update to {len(self.active_connections)} clients...")
+        # Create a list of tasks to send updates concurrently
+        # Make a copy of the set for safe iteration during potential disconnects
+        send_tasks = [
+             connection.send_text(html_fragment)
+             for connection in list(self.active_connections) # Iterate over a copy
+        ]
+        # Use asyncio.gather to send to all clients, collecting results/exceptions
+        results = await asyncio.gather(*send_tasks, return_exceptions=True)
+        # Handle potential errors during broadcast (e.g., client disconnected abruptly)
+        # Iterate over connections again, checking results
+        connections_to_remove = set()
+        for i, result in enumerate(results):
+            connection = list(self.active_connections)[i] # Assumes order is maintained
+            if isinstance(result, Exception):
+                print(f"Error sending update to client: {result}. Marking for removal.")
+                connections_to_remove.add(connection)
+        # Disconnect clients that failed
+        for connection in connections_to_remove:
+            await self.disconnect(connection)
 manager = ConnectionManager()
 # --- API Routes ---
 @app.get("/", response_class=HTMLResponse)
 async def get_homepage():
+    """Serves the main HTML page with WebSocket connection and improved styling."""
+    # NOTE: Ensure the static path '/static/pokemon_huggingface.png' is correct
+    # and the image exists in a 'static' folder next to your main.py
     return """
     <!DOCTYPE html>
+    <html lang="en">
     <head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
         <title>Pokemon Battle Livestream</title>
+        <link rel="preconnect" href="https://fonts.googleapis.com">
+        <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+        <link href="https://fonts.googleapis.com/css2?family=Poppins:wght@400;700&family=Press+Start+2P&display=swap" rel="stylesheet">
         <style>
+            /* Basic Reset */
+            * {
+                box-sizing: border-box;
+            }
+            html, body {
                 margin: 0;
                 padding: 0;
                 height: 100%;
+                width: 100%;
+                overflow: hidden; /* Prevent scrollbars on body */
+                font-family: 'Poppins', sans-serif; /* Default font */
+                color: #ffffff; /* Default text color */
+                background-color: #1a1a1a; /* Dark background */
             }
+            /* Container for dynamic content */
             #stream-container {
+                position: fixed; /* Use fixed to ensure it covers viewport */
                 top: 0;
                 left: 0;
                 width: 100%;
                 height: 100%;
+                display: flex; /* Use flexbox for centering content */
+                justify-content: center;
+                align-items: center;
+            }
+            /* Iframe Styling */
+            .battle-iframe {
+                width: 100%;
+                height: 100%;
+                border: none; /* Remove default border */
+                display: block; /* Prevents potential extra space below iframe */
+            }
+            /* Base Content Container Styling (used by idle/error) */
+            .content-container {
+                width: 100%;
+                height: 100%;
+                display: flex;
+                flex-direction: column;
+                justify-content: center;
+                align-items: center;
+                padding: 20px;
+                text-align: center;
+            }
+            /* Idle Screen Specific Styling */
+            .idle-container {
+                 /* Ensure the background covers the entire container */
+                background-image: url('/static/pokemon_huggingface.png');
+                background-size: cover;
+                background-position: center;
+                background-repeat: no-repeat;
+            }
+            /* Error Screen Specific Styling */
+            .error-container {
+                background: linear-gradient(135deg, #4d0000, #1a0000); /* Dark red gradient */
+            }
+            /* Message Box Styling (shared by idle/error) */
+            .message-box {
+                background-color: rgba(0, 0, 0, 0.75); /* Darker, more opaque */
+                padding: 40px 50px; /* More padding */
+                border-radius: 20px; /* More rounded */
+                max-width: 70%; /* Max width */
+                box-shadow: 0 8px 25px rgba(0, 0, 0, 0.5); /* Softer shadow */
+                border: 1px solid rgba(255, 255, 255, 0.1); /* Subtle border */
+            }
+            .status {
+                font-family: 'Press Start 2P', cursive; /* Pixel font for status */
+                font-size: clamp(1.5em, 4vw, 2.5em); /* Responsive font size */
+                margin-bottom: 25px;
+                color: #ffcb05; /* Pokemon Yellow */
+                text-shadow: 3px 3px 0px #3b4cca; /* Pokemon Blue shadow */
+                /* Subtle pulse animation for idle status */
+                animation: pulse 2s infinite ease-in-out;
+            }
+            .instruction {
+                font-size: clamp(1em, 2.5vw, 1.4em); /* Responsive font size */
+                color: #f0f0f0; /* Light grey for readability */
+                line-height: 1.6;
+                text-shadow: 1px 1px 3px rgba(0, 0, 0, 0.7);
+            }
+             .instruction strong {
+                 color: #ff7f0f; /* A contrasting color like orange */
+                 font-weight: 700; /* Ensure Poppins bold is used */
+             }
+            /* Error Screen Specific Text Styling */
+            .error-container .status {
+                color: #ff4d4d; /* Bright Red for error status */
+                text-shadow: 2px 2px 0px #800000; /* Darker red shadow */
+                animation: none; /* No pulse on error */
+            }
+             .error-container .instruction {
+                color: #ffdddd; /* Lighter red for error details */
+            }
+            /* Pulse Animation */
+            @keyframes pulse {
+                0% { transform: scale(1); }
+                50% { transform: scale(1.03); }
+                100% { transform: scale(1); }
             }
         </style>
     </head>
     <body>
+        <div id="stream-container">
+             </div>
         <script>
             const streamContainer = document.getElementById('stream-container');
+            let ws = null; // WebSocket instance
+            function connectWebSocket() {
+                // Use wss:// for https:// and ws:// for http://
+                const wsProtocol = location.protocol === 'https:' ? 'wss' : 'ws';
+                const wsUrl = `${wsProtocol}://${location.host}/ws`;
+                ws = new WebSocket(wsUrl);
+                console.log('Attempting to connect to WebSocket server...');
+                ws.onopen = (event) => {
+                    console.log('WebSocket connection established.');
+                    // Optional: Clear any 'connecting...' message if you have one
+                    // streamContainer.innerHTML = ''; // Clear container only if needed
+                };
+                ws.onmessage = (event) => {
+                    // console.log('Received update from server:', event.data);
+                    // Directly set the innerHTML with the fragment received from the server
+                    streamContainer.innerHTML = event.data;
+                };
+                ws.onclose = (event) => {
+                    console.log(`WebSocket connection closed. Code: ${event.code}, Reason: ${event.reason}. Attempting to reconnect in 5 seconds...`);
+                    ws = null; // Clear the instance
+                    // Clear the display or show a 'disconnected' message
+                    streamContainer.innerHTML = createReconnectMessage();
+                    setTimeout(connectWebSocket, 5000); // Retry connection after 5 seconds
+                };
+                ws.onerror = (event) => {
+                    console.error('WebSocket error:', event);
+                    // The onclose event will usually fire after an error,
+                    // so reconnection logic is handled there.
+                    // You might want to display an error message here briefly.
+                    streamContainer.innerHTML = createErrorMessage("WebSocket connection error. Attempting to reconnect...");
+                    // Optionally force close to trigger reconnect logic if onclose doesn't fire
+                    if (ws && ws.readyState !== WebSocket.CLOSED) {
+                        ws.close();
+                    }
+                };
+            }
+             // Helper function to generate reconnecting message HTML (matches error style)
+            function createReconnectMessage() {
+                 return `
+                    <div class="content-container error-container" style="background: #333;">
+                        <div class="message-box" style="background-color: rgba(0,0,0,0.6);">
+                            <p class="status" style="color: #ffcb05; text-shadow: none; animation: none;">🔌 Disconnected 🔌</p>
+                            <p class="instruction" style="color: #eee;">Connection lost. Attempting to reconnect automatically...</p>
+                        </div>
+                    </div>`;
+            }
+             // Helper function to generate error message HTML
+            function createErrorMessage(message) {
+                return `
+                    <div class="content-container error-container">
+                        <div class="message-box">
+                            <p class="status">🚨 Error 🚨</p>
+                            <p class="instruction">${message}</p>
+                        </div>
+                    </div>`;
+            }
+            // Initial connection attempt when the page loads
+            connectWebSocket();
         </script>
     </body>
     </html>
     await manager.connect(websocket)
     try:
         while True:
+            # Keep connection alive. Client doesn't send messages in this setup.
+            # FastAPI's WebSocket implementation handles ping/pong internally usually.
+            # If needed, you could implement explicit keepalive here.
+            data = await websocket.receive_text()
+            # We don't expect messages from the client in this design,
+            # but log if received for debugging.
+            print(f"Received unexpected message from client: {data}")
+            # Or simply keep listening:
+            # await asyncio.sleep(60) # Example keepalive interval if needed
+    except WebSocketDisconnect as e:
+         print(f"WebSocket disconnected: Code {e.code}, Reason: {getattr(e, 'reason', 'N/A')}")
+         await manager.disconnect(websocket) # Use await here
     except Exception as e:
+        # Catch other potential errors on the connection
         print(f"WebSocket error: {e}")
+        traceback.print_exc()
+        await manager.disconnect(websocket) # Ensure disconnect on error
 @app.on_event("startup")
 async def startup_event():
     """Start background tasks when the application starts."""
     global background_task_handle
+    # Mount static files directory (make sure 'static' folder exists)
+    # Place your 'pokemon_huggingface.png' inside this 'static' folder
+    static_dir = "static"
+    if not os.path.exists(static_dir):
+        os.makedirs(static_dir)
+        print(f"Created static directory at: {os.path.abspath(static_dir)}")
+        print("!!! Please add 'pokemon_huggingface.png' to this directory! !!!")
+    app.mount("/static", StaticFiles(directory=static_dir), name="static")
+    print(f"Mounted static directory '{static_dir}' at '/static'")
     print("🚀 Starting background tasks")
+    # Start the main lifecycle manager task
+    background_task_handle = asyncio.create_task(manage_agent_lifecycle(), name="LifecycleManager")
+    # Add the exception logging callback
     background_task_handle.add_done_callback(log_task_exception)
     print("✅ Background tasks started")
 async def shutdown_event():
     """Clean up tasks when shutting down."""
     global background_task_handle, active_agent_instance
+    print("\n🔌 Shutting down application. Cleaning up...")
+    # 1. Cancel the main lifecycle manager task
     if background_task_handle and not background_task_handle.done():
+        print("Cancelling background task...")
         background_task_handle.cancel()
         try:
             await asyncio.wait_for(background_task_handle, timeout=5.0)
+            print("Background task cancelled successfully.")
+        except asyncio.CancelledError:
+            print("Background task cancellation confirmed (CancelledError).")
+        except asyncio.TimeoutError:
+            print("Background task did not finish cancelling within timeout.")
+        except Exception as e:
+            print(f"Error during background task cancellation: {e}")
+    # 2. Deactivate and disconnect any currently active agent
+    #    Use a copy of the instance in case it gets cleared elsewhere during shutdown.
+    agent_to_disconnect = active_agent_instance
+    if agent_to_disconnect:
+        agent_name = agent_to_disconnect.username if hasattr(agent_to_disconnect, 'username') else 'Unknown Agent'
+        print(f"Disconnecting active agent '{agent_name}'...")
         try:
+             # Check websocket status before disconnecting
+             if hasattr(agent_to_disconnect, '_websocket') and agent_to_disconnect._websocket and agent_to_disconnect._websocket.open:
+                 await agent_to_disconnect.disconnect()
+                 print(f"Agent '{agent_name}' disconnected.")
+             else:
+                 print(f"Agent '{agent_name}' already disconnected or websocket not available.")
         except Exception as e:
+            print(f"Error during agent disconnect on shutdown for '{agent_name}': {e}")
+    # 3. Close all active WebSocket connections cleanly
+    print(f"Closing {len(manager.active_connections)} client WebSocket connections...")
+    # Create tasks to close all connections concurrently
+    close_tasks = [
+         conn.close(code=1000, reason="Server shutting down") # 1000 = Normal Closure
+         for conn in list(manager.active_connections) # Iterate over a copy
+    ]
+    if close_tasks:
+        await asyncio.gather(*close_tasks, return_exceptions=True) # Log potential errors during close
+    print("✅ Cleanup complete. Application shutdown.")
 # For direct script execution
 if __name__ == "__main__":
     import uvicorn
+    # Configure logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s:%(lineno)d - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    # Reduce noise from poke_env's default INFO logging if desired
     logging.getLogger('poke_env').setLevel(logging.WARNING)
+    logging.getLogger('websockets.client').setLevel(logging.INFO) # Show websocket connection attempts
     print("Starting Pokemon Battle Livestream Server...")
+    print("="*60)
     if not AVAILABLE_AGENT_NAMES:
+        print("█████████████████████ FATAL ERROR █████████████████████")
+        print(" No agents found with configured passwords!")
+        print(" Please set the required environment variables:")
         for name, cfg in AGENT_CONFIGS.items():
+            print(f"  - {cfg.get('password_env_var', 'N/A')} (for agent: {name})")
+        print("="*60)
         exit("Exiting due to missing agent passwords.")
     else:
+        print("✨ Available Agents Found:")
         for name in AVAILABLE_AGENT_NAMES:
+            print(f"  - {name}")
+    print("="*60)
+    print(f"Server will run on http://0.0.0.0:7860")
+    print("="*60)
     # Run with uvicorn
+    uvicorn.run(
+        "main:app", # Point to the FastAPI app instance
+        host="0.0.0.0",
+        port=7860,
+        reload=False, # Disable reload for production/stable testing
+        log_level="info" # Uvicorn's log level
+    )