Update main.py

main.py

@@ -6,18 +6,14 @@ import random
 import time
 import traceback
 import logging
-# --- Additions for last_action route ---
-import datetime
-import html
-from typing import List, Dict, Optional, Set, Callable # Added Callable
-# ---------------------------------------
+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, ActionType, ForfeitAction, MoveOrder, SwitchOrder, DefaultOrder # Import action types
+from poke_env.player import Player
 from poke_env import AccountConfiguration, ServerConfiguration
 from poke_env.environment.battle import Battle
 
@@ -54,26 +50,10 @@ 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
-# --- NEW: Global variable for last action ---
-last_llm_action: Optional[Dict] = None
-# --------------------------------------------
 
 # --- Create FastAPI app ---
 app = FastAPI(title="Pokemon Battle Livestream")
 
-
-# --- NEW: Callback function for agents ---
-def update_last_action_callback(action_info: Dict):
-    """Callback for agents to report their chosen action."""
-    global last_llm_action
-    # Add a timestamp for context
-    action_info["timestamp_utc"] = datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S UTC')
-    last_llm_action = action_info
-    # Optional: Log the update for debugging
-    print(f"ACTION_LOG: Agent '{action_info.get('agent', 'Unknown')}' chose action: {action_info.get('action_str', 'N/A')} (Turn: {action_info.get('turn', '?')})")
-# ---------------------------------------
-
-
 # --- Helper Functions ---
 def get_active_battle(agent: Player) -> Optional[Battle]:
     """Returns the first non-finished battle for an agent."""
@@ -97,7 +77,11 @@ def get_active_battle(agent: Player) -> Optional[Battle]:
 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"
@@ -109,6 +93,7 @@ def create_battle_iframe(battle_id: str) -> str:
 
 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">
@@ -120,6 +105,7 @@ def create_idle_html(status_message: str, instruction: str) -> str:
 
 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">
@@ -129,139 +115,9 @@ def create_error_html(error_msg: str) -> str:
     </div>
     """
 
-# --- NEW: Helper function to create HTML for the last action page ---
-def create_last_action_html(action_data: Optional[Dict]) -> str:
-    """Formats the last action data into an HTML page."""
-    if not action_data:
-        content = """
-            <div class="message-box">
-                <p class="status">No Action Recorded Yet</p>
-                <p class="instruction">Waiting for the agent to make its first move...</p>
-            </div>
-        """
-    else:
-        # Escape HTML characters in potentially user-generated content
-        raw_output_escaped = html.escape(action_data.get("raw_llm_output", "N/A"))
-        agent_name_escaped = html.escape(action_data.get('agent', 'Unknown'))
-        action_type_escaped = html.escape(action_data.get('action_type', 'N/A'))
-        action_str_escaped = html.escape(action_data.get('action_str', 'N/A'))
-
-        content = f"""
-            <div class="action-details">
-                <h2>Last Action by Agent: <span class="agent-name">{agent_name_escaped}</span></h2>
-                <p><strong>Timestamp:</strong> {action_data.get('timestamp_utc', 'N/A')}</p>
-                <p><strong>Battle Turn:</strong> {action_data.get('turn', '?')}</p>
-                <p><strong>Action Type:</strong> {action_type_escaped}</p>
-                <p><strong>Action Chosen:</strong> <span class="action-chosen">{action_str_escaped}</span></p>
-                <div class="raw-output">
-                    <h3>Raw LLM Output (if available):</h3>
-                    <pre>{raw_output_escaped}</pre>
-                </div>
-            </div>
-        """
-
-    # Basic HTML structure with auto-refresh and styling
-    return f"""
-    <!DOCTYPE html>
-    <html lang="en">
-    <head>
-        <meta charset="UTF-8">
-        <meta name="viewport" content="width=device-width, initial-scale=1.0">
-        <meta http-equiv="refresh" content="5">
-        <title>Last Agent Action</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>
-            body {{
-                font-family: 'Poppins', sans-serif;
-                background-color: #2c2f33;
-                color: #ffffff;
-                margin: 0;
-                padding: 20px;
-                line-height: 1.6;
-            }}
-            .container {{
-                max-width: 900px;
-                margin: 20px auto;
-                background-color: #3e4147;
-                padding: 25px;
-                border-radius: 10px;
-                box-shadow: 0 4px 15px rgba(0, 0, 0, 0.3);
-            }}
-             h1, h2 {{
-                 font-family: 'Press Start 2P', cursive;
-                 color: #ffcb05; /* Pokemon Yellow */
-                 text-shadow: 2px 2px 0px #3b4cca; /* Pokemon Blue shadow */
-                 margin-bottom: 20px;
-                 text-align: center;
-             }}
-            .action-details p {{
-                margin: 10px 0;
-                font-size: 1.1em;
-            }}
-            .action-details strong {{
-                color: #ffcb05; /* Pokemon Yellow */
-            }}
-            .agent-name {{
-                color: #ff7f0f; /* Orange */
-                font-weight: bold;
-            }}
-            .action-chosen {{
-                font-weight: bold;
-                color: #76d7c4; /* Teal */
-                font-size: 1.2em;
-            }}
-            .raw-output {{
-                margin-top: 25px;
-                border-top: 1px solid #555;
-                padding-top: 15px;
-            }}
-             .raw-output h3 {{
-                 margin-bottom: 10px;
-                 color: #f0f0f0;
-             }}
-            pre {{
-                background-color: #23272a;
-                color: #dcdcdc;
-                padding: 15px;
-                border-radius: 5px;
-                white-space: pre-wrap; /* Wrap long lines */
-                word-wrap: break-word; /* Break words if necessary */
-                font-family: monospace;
-                font-size: 0.95em;
-                max-height: 300px; /* Limit height */
-                overflow-y: auto;  /* Add scrollbar if needed */
-            }}
-             /* Styling for the 'No Action' message box */
-            .message-box {{
-                text-align: center;
-                padding: 30px;
-            }}
-            .message-box .status {{
-                 font-family: 'Press Start 2P', cursive;
-                 font-size: 1.8em;
-                 color: #ff7f0f; /* Orange */
-                 margin-bottom: 15px;
-            }}
-            .message-box .instruction {{
-                font-size: 1.1em;
-                color: #cccccc;
-            }}
-        </style>
-    </head>
-    <body>
-        <div class="container">
-            <h1>Agent Action Log</h1>
-            {content}
-        </div>
-    </body>
-    </html>
-    """
-# --------------------------------------------------------------------
-
 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.")
 
@@ -269,13 +125,7 @@ async def update_display_html(new_html_fragment: str) -> None:
 # --- Agent Lifecycle Management ---
 async def select_and_activate_new_agent():
     """Selects a random available agent, instantiates it, and starts its listening task."""
-    # --- MODIFIED: Make sure globals are declared ---
-    global active_agent_name, active_agent_instance, active_agent_task, last_llm_action
-    # -----------------------------------------------
-
-    # --- MODIFIED: Reset last action when selecting a new agent ---
-    last_llm_action = None
-    # -----------------------------------------------------------
+    global active_agent_name, active_agent_instance, active_agent_task
 
     if not AVAILABLE_AGENT_NAMES:
         print("Lifecycle: No available agents with passwords set.")
@@ -289,29 +139,31 @@ async def select_and_activate_new_agent():
     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)
-        # --- MODIFIED: Pass the action callback to the agent constructor ---
         agent = AgentClass(
             account_configuration=account_config,
             server_configuration=custom_config,
             battle_format=DEFAULT_BATTLE_FORMAT,
             log_level=logging.INFO,
-            max_concurrent_battles=1,
-            action_callback=update_last_action_callback # Pass the callback function
+            max_concurrent_battles=1
         )
-        # -------------------------------------------------------------------
 
+        # 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)
+        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
@@ -322,6 +174,7 @@ async def select_and_activate_new_agent():
         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
         active_agent_task = None
@@ -329,24 +182,38 @@ async def select_and_activate_new_agent():
 
 async def check_for_new_battle():
     """Checks if the active agent has started a battle with a valid tag."""
+    # --- FIX: Declare intention to use/modify global variables ---
     global active_agent_instance, current_battle_instance, active_agent_name, active_agent_task
+    # -------------------------------------------------------------
 
     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:
+            # This line MODIFIES the global variable
             current_battle_instance = battle
             print(f"Lifecycle: Agent '{active_agent_name}' started battle: {battle.battle_tag}")
+
+            # 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
+    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":
@@ -356,22 +223,35 @@ async def deactivate_current_agent(reason: str = "cycle"):
     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>..."))
 
-    await asyncio.sleep(3)
+    # 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:
@@ -379,32 +259,39 @@ async def deactivate_current_agent(reason: str = "cycle"):
         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()
+            traceback.print_exc() # Log full traceback for debugging
 
-    await asyncio.sleep(2)
+    # 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.")
 
 async def manage_agent_lifecycle():
     """Runs the main loop selecting, running, and cleaning up agents sequentially."""
+    # --- FIX: Declare intention to use global variables ---
     global active_agent_name, active_agent_instance, active_agent_task, current_battle_instance
+    # ------------------------------------------------------
 
     print("Background lifecycle manager started.")
-    REFRESH_INTERVAL_SECONDS = 3
-    LOOP_COOLDOWN_SECONDS = 1
-    ERROR_RETRY_DELAY_SECONDS = 10
-    POST_BATTLE_DELAY_SECONDS = 5
+    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
 
@@ -414,6 +301,10 @@ async def manage_agent_lifecycle():
         print(f"\n--- Lifecycle Check #{loop_counter} [{time.strftime('%H:%M:%S')}] ---")
 
         try:
+            # ==================================
+            # State 1: No agent active
+            # ==================================
+            # Now Python knows active_agent_instance refers to the global one
             if active_agent_instance is None:
                 print(f"[{loop_counter}] State 1: No active agent. Selecting...")
                 activated = await select_and_activate_new_agent()
@@ -421,52 +312,82 @@ async def manage_agent_lifecycle():
                     print(f"[{loop_counter}] State 1: Activation failed. Waiting {ERROR_RETRY_DELAY_SECONDS}s before retry.")
                     await asyncio.sleep(ERROR_RETRY_DELAY_SECONDS)
                 else:
+                    # Now Python knows active_agent_name refers to the global one set by select_and_activate_new_agent
                     print(f"[{loop_counter}] State 1: Agent '{active_agent_name}' activated successfully.")
+                    # No sleep here, proceed to next check immediately if needed
 
-            else: # Agent is active
-                agent_name = active_agent_name
+            # ==================================
+            # State 2: Agent is active
+            # ==================================
+            else:
+                # Now Python knows active_agent_name refers to the global one
+                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 ---
+                # Now Python knows current_battle_instance refers to the global one
                 if current_battle_instance is None:
                     print(f"[{loop_counter}] State 2a: Checking for new battle for '{agent_name}'...")
-                    await check_for_new_battle()
+                    await check_for_new_battle() # This updates global current_battle_instance if found
 
+                    # Now Python knows current_battle_instance refers to the global one
                     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:
+                                # Now Python knows active_agent_instance refers to the global one
+                                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)
+                                    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
+                            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)
+                        await asyncio.sleep(REFRESH_INTERVAL_SECONDS) # Wait before next check if idle
+
 
-                if current_battle_instance is not None: # Check again in case a battle just started
+                # --- Sub-state: Monitor ongoing battle ---
+                 # Now Python knows current_battle_instance refers to the global one
+                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
+                    # Now Python knows active_agent_instance refers to the global one
                     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
+                    # Now Python knows active_agent_instance refers to the global one
                     battle_obj = active_agent_instance._battles.get(battle_tag)
 
                     if battle_obj and battle_obj.finished:
@@ -474,109 +395,138 @@ async def manage_agent_lifecycle():
                         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
+                        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}'.")
-                        await asyncio.sleep(REFRESH_INTERVAL_SECONDS)
+                        # 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
+            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_err = active_agent_name # Use different var name to avoid conflict
+             # Now Python knows active_agent_name refers to the global one
+            current_agent_name = active_agent_name # Cache name before deactivation attempts
+            # Now Python knows active_agent_instance refers to the global one
             if active_agent_instance:
-                print(f"Attempting to deactivate agent '{current_agent_name_err}' due to loop error...")
+                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
+            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
+         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:
-             return
+             # 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)
+             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()
-        # Need to iterate carefully if connections can change during gather
-        conn_list = list(self.active_connections)
         for i, result in enumerate(results):
-             # Ensure index is valid if connections changed mid-gather
-             if i < len(conn_list):
-                 connection = conn_list[i]
-                 if isinstance(result, Exception):
-                     print(f"Error sending update to client: {result}. Marking for removal.")
-                     connections_to_remove.add(connection)
+            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:
-            # Check if connection still exists before disconnecting
-            if connection in self.active_connections:
-                 await self.disconnect(connection)
+            await self.disconnect(connection)
 
 
 manager = ConnectionManager()
@@ -585,7 +535,8 @@ manager = ConnectionManager()
 @app.get("/", response_class=HTMLResponse)
 async def get_homepage():
     """Serves the main HTML page with WebSocket connection and improved styling."""
-    # ... (HTML remains the same as before)
+    # 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">
@@ -601,7 +552,6 @@ async def get_homepage():
             * {
                 box-sizing: border-box;
             }
-
             html, body {
                 margin: 0;
                 padding: 0;
@@ -612,7 +562,6 @@ async def get_homepage():
                 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 */
@@ -624,7 +573,6 @@ async def get_homepage():
                 justify-content: center;
                 align-items: center;
             }
-
             /* Iframe Styling */
             .battle-iframe {
                 width: 100%;
@@ -632,7 +580,6 @@ async def get_homepage():
                 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%;
@@ -644,7 +591,6 @@ async def get_homepage():
                 padding: 20px;
                 text-align: center;
             }
-
             /* Idle Screen Specific Styling */
             .idle-container {
                  /* Ensure the background covers the entire container */
@@ -653,12 +599,10 @@ async def get_homepage():
                 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 */
@@ -668,7 +612,6 @@ async def get_homepage():
                 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 */
@@ -678,7 +621,6 @@ async def get_homepage():
                 /* 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 */
@@ -689,8 +631,6 @@ async def get_homepage():
                  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 */
@@ -700,56 +640,56 @@ async def get_homepage():
              .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;
+                    ws = null; // Clear the instance
+                    // Clear the display or show a 'disconnected' message
                     streamContainer.innerHTML = createReconnectMessage();
-                    setTimeout(connectWebSocket, 5000);
+                    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;">
@@ -759,7 +699,7 @@ async def get_homepage():
                         </div>
                     </div>`;
             }
-
+             // Helper function to generate error message HTML
             function createErrorMessage(message) {
                 return `
                     <div class="content-container error-container">
@@ -769,7 +709,7 @@ async def get_homepage():
                         </div>
                     </div>`;
             }
-
+            // Initial connection attempt when the page loads
             connectWebSocket();
         </script>
     </body>
@@ -781,31 +721,32 @@ async def websocket_endpoint(websocket: WebSocket):
     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)
+         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)
-
-# --- NEW: Route to display the last action ---
-@app.get("/last_action", response_class=HTMLResponse)
-async def get_last_action():
-    """Serves an HTML page displaying the last recorded agent action."""
-    global last_llm_action
-    # Return the formatted HTML page using the current state
-    return create_last_action_html(last_llm_action)
-# --------------------------------------------
-
-# --- Lifecyle Events ---
+        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)
@@ -816,7 +757,9 @@ async def startup_event():
     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")
 
@@ -827,6 +770,7 @@ async def shutdown_event():
 
     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()
@@ -840,11 +784,14 @@ async def shutdown_event():
         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.")
@@ -853,28 +800,32 @@ async def shutdown_event():
         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")
-         for conn in list(manager.active_connections)
+         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)
+        await asyncio.gather(*close_tasks, return_exceptions=True) # Log potential errors during close
 
     print("✅ Cleanup complete. Application shutdown.")
 
 
-# --- Main execution ---
+# 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)
+    logging.getLogger('websockets.client').setLevel(logging.INFO) # Show websocket connection attempts
 
     print("Starting Pokemon Battle Livestream Server...")
     print("="*60)
@@ -893,13 +844,14 @@ if __name__ == "__main__":
             print(f"  - {name}")
     print("="*60)
     print(f"Server will run on http://0.0.0.0:7860")
-    print("Access the action log at http://0.0.0.0:7860/last_action") # Added info
     print("="*60)
 
+    # Run with uvicorn
     uvicorn.run(
-        "main:app",
+        "main:app", # Point to the FastAPI app instance
         host="0.0.0.0",
         port=7860,
-        reload=False,
-        log_level="info"
-    )
+        reload=False, # Disable reload for production/stable testing
+        log_level="info" # Uvicorn's log level
+    )
+