from pprint import pprint
import os
import signal
import sys
from orchestrator_maze_implementation.utils.maze_utils import create_maze_by_uuid
from orchestrator_maze_implementation.visualization import display_maze_state
from orchestrator_maze_implementation.visualization.maze_visualization import get_agent_offset_x
from orchestrator_maze_implementation.utils.summarize_agent_conversation import display_conversation_summary
from orchestrator_maze_implementation.state import MazeState
from orchestrator_maze_implementation.router.maze_router import create_maze_workflow
from orchestrator_maze_implementation.router.router_support import get_current_agent_by_turn
from orchestrator_maze_implementation.config.config_service import ConfigService
from copy import deepcopy
import matplotlib.pyplot as plt
import matplotlib.patches
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.runnables import chain
import pickle
from datetime import datetime
import argparse

# Environment variable loading
try:
    from dotenv import load_dotenv
    DOTENV_AVAILABLE = True
except ImportError:
    DOTENV_AVAILABLE = False
    print("⚠️  python-dotenv not available. Environment variables will not be loaded from .env file.")
    print("   Install with: pip install python-dotenv")

# Configuration file path - can be changed to point to different config files
# Default config.yaml relative to this script's location
def get_default_config_path():
    """Get the default config file path relative to this script's location"""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(script_dir, 'config.yaml')

CONFIG_FILE_PATH = get_default_config_path()

# Global variables for graceful shutdown
_current_state = None
_config_service = None 
_maze_uuid = None
_config_path = None
_output_dir = None
_graceful_shutdown_initiated = False
_final_outputs_generated = False  # Track whether final outputs were generated

def _save_emergency_state(state, config_service, maze_uuid, config_path, output_dir):
    """
    Emergency state saving function for graceful shutdown scenarios.
    This ensures scientific reproducibility even when processes are terminated.
    """
    global _graceful_shutdown_initiated
    
    if _graceful_shutdown_initiated:
        return  # Prevent multiple invocations
    
    _graceful_shutdown_initiated = True
    
    try:
        print("\n🚨 EMERGENCY: Saving final state for reproducibility...")
        
        # Call the existing generate_final_outputs function
        generate_final_outputs(
            state, 
            config_service, 
            maze_uuid, 
            config_path, 
            output_dir, 
            state_was_recovered=True
        )
        
        print("✅ Emergency state save completed")
        
    except Exception as e:
        print(f"❌ Emergency state save failed: {e}")
        
        # Last resort: try to save at least the pickle file
        try:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            emergency_pickle = os.path.join(output_dir, f"emergency_state_{maze_uuid}_{timestamp}.pickle")
            
            with open(emergency_pickle, 'wb') as f:
                pickle.dump({
                    "final_state": state,
                    "config": config_service.get_config() if config_service else None,
                    "maze_uuid": maze_uuid,
                    "timestamp": timestamp,
                    "emergency_save": True,
                    "error_during_save": str(e)
                }, f, protocol=pickle.HIGHEST_PROTOCOL)
            
            print(f"✅ Emergency pickle saved: {emergency_pickle}")
            
        except Exception as pickle_error:
            print(f"❌ Even emergency pickle save failed: {pickle_error}")

def _signal_handler(signum, frame):
    """Signal handler for graceful shutdown"""
    global _current_state, _config_service, _maze_uuid, _config_path, _output_dir
    
    print(f"\n🚨 Received signal {signum} - initiating graceful shutdown...")
    
    if _current_state is not None:
        _save_emergency_state(_current_state, _config_service, _maze_uuid, _config_path, _output_dir)
    else:
        print("⚠️  No current state available for emergency save")
    
    # Exit with appropriate code
    sys.exit(128 + signum)

# Register signal handlers for graceful shutdown
signal.signal(signal.SIGTERM, _signal_handler)
signal.signal(signal.SIGINT, _signal_handler)  # Ctrl+C

def _get_safe_mode_from_config_or_env(config_service: ConfigService) -> bool:
    # Single source of truth: config service handles env/YAML/auto-detect
    return config_service.get_visualization_safe_mode()


def main(config_path: str = None, output_dir: str = None):
    """Main function to run the maze navigation

    Args:
        config_path: Optional path to a YAML config. If not provided, uses default.
        output_dir: Optional output directory root where final artifacts (images, pickles)
            will be saved. If not provided, falls back to `<config_dir>/outputs`.
    """
    global _current_state, _config_service, _maze_uuid, _config_path, _output_dir
    
    # Load environment variables from .env file relative to this script's location
    if DOTENV_AVAILABLE:
        # Get the directory where this script is located
        script_dir = os.path.dirname(os.path.abspath(__file__))
        env_file_path = os.path.join(script_dir, '.env')
        
        # Check if .env file exists and load it
        if os.path.exists(env_file_path):
            load_dotenv(env_file_path)
            print(f"✅ Environment variables loaded from: {env_file_path}")
        else:
            print(f"ℹ️  No .env file found at: {env_file_path}")
            print("   Create a .env file in the same directory as main.py to load environment variables")
            raise Exception("No .env file found")
    
    # Use provided config path or default
    if config_path is None:
        ## it is a lie, it can be reached
        config_path = CONFIG_FILE_PATH
    print(f"Using config path: {config_path}")
    
    # Set global variables for signal handling
    _config_path = config_path
    
    # Resolve output directory root for artifacts
    # If not provided, default to sibling `outputs/` next to the EFFECTIVE config file
    try:
        if output_dir is not None:
            RESOLVED_OUTPUT_DIR = os.path.abspath(output_dir)
        else:
            config_dir_default = os.path.dirname(os.path.abspath(config_path))
            RESOLVED_OUTPUT_DIR = os.path.join(config_dir_default, "outputs")
        os.makedirs(RESOLVED_OUTPUT_DIR, exist_ok=True)
        print(f"Artifacts output directory: {RESOLVED_OUTPUT_DIR}")
        
        # Set global variable
        _output_dir = RESOLVED_OUTPUT_DIR
        
    except Exception as e:
        print(f"❌ Failed to prepare output directory '{output_dir}': {e}")
        return
        
    # Initialize the ConfigService singleton
    config_service = ConfigService()
    
    try:
        config_service.load_config(config_path)
        print("✅ Configuration loaded successfully from:", config_path)
        
        # Set global variables
        _config_service = config_service
        
    except (FileNotFoundError, Exception) as e:
        print(f"❌ Error loading configuration: {e}")
        return
    
    
    print("🧩 Multi-Agent Maze Solver Starting...")
    print(" Experiment Settings are: ")
    pprint(config_service.get_config())
    print("=" * 50)
    
    # Use config service methods for initialization
    num_agents = config_service.get_num_execution_agents()
    maze_uuid = config_service.get_maze_uuid()
    
    # Set global variables
    _maze_uuid = maze_uuid
    
    state = initialize_state(num_agents=num_agents, maze_uuid=maze_uuid)
    
    # Set global state for signal handling
    _current_state = state

    # Create workflow
    workflow = create_maze_workflow(state)
    
    # Run the maze navigation
    print("\nStarting continuous maze navigation...")
    print("The agent will continuously navigate with 0.5 second delays between moves")
    print("The agent will use ToolNode for automatic tool handling")
    print("Maze state will be displayed using matplotlib visualizations")
    print("Navigation will continue until the maze is solved or max turns reached")
    
    steps_per_agent = config_service.get_steps_per_agent()
    print(f"🔄 {len(state['all_agents'])} agents will rotate every {steps_per_agent} steps: {' -> '.join(state['all_agents'])} -> {state['all_agents'][0]} (cycle)")
    print(f"🎯 Using maze UUID: {maze_uuid}")
    print("=" * 50)
    
    try:
        # Use config service to get recursion limit
        recursion_limit = config_service.get_recursion_limit()
        workflow_config = {"recursion_limit": recursion_limit}
        final_state = workflow.invoke(state, config=workflow_config)
        
        # Update global state
        _current_state = final_state
        
        print("\n" + "=" * 50)
        print("MAZE NAVIGATION COMPLETE!")
        print("=" * 50)
        
        # Generate final outputs (visualization, pickle, etc.)
        generate_final_outputs(final_state, config_service, maze_uuid, config_path, RESOLVED_OUTPUT_DIR, state_was_recovered=False)
        
        # Mark that final outputs were generated
        global _final_outputs_generated
        _final_outputs_generated = True
        
    except Exception as e:
        print(f"Error during maze navigation: {e}")
        import traceback
        traceback.print_exc()
        
        # Attempt to recover final state from workflow checkpointer on timeout/error
        final_state = None
        try:
            print("\n⚠️  Attempting to recover final state from workflow checkpointer...")
            
            # Try to get the last known state from the workflow's checkpointer
            # The workflow config should have thread_id or similar identifier
            if hasattr(workflow, 'get_state'):
                # Create a config that matches the workflow execution
                recovery_config = {"configurable": {"thread_id": "default"}}
                
                # Get the last checkpoint state
                last_checkpoint = workflow.get_state(recovery_config)
                if last_checkpoint and hasattr(last_checkpoint, 'values') and last_checkpoint.values:
                    final_state = last_checkpoint.values
                    print(f"✅ Recovered final state from checkpointer with {len(final_state.get('messages', []))} messages")
                    print(f"   Turn count: {final_state.get('turn_count', 'unknown')} turns")
                else:
                    print("❌ No recoverable state found in checkpointer")
            else:
                print("❌ Workflow does not support state recovery")
                
        except Exception as recovery_error:
            print(f"❌ Failed to recover state from checkpointer: {recovery_error}")
        
        # Use current state if no recovery possible
        if final_state is None and _current_state is not None:
            final_state = _current_state
            print("⚠️  Using last known state from global variable")
        
        # If we have any state available, proceed with visualization and output generation
        if final_state:
            # Update global state
            _current_state = final_state
            
            print("\n" + "=" * 50)
            print("MAZE NAVIGATION INTERRUPTED - PARTIAL RESULTS RECOVERED")
            print("=" * 50)
            
            # Generate final outputs with recovered state
            generate_final_outputs(final_state, config_service, maze_uuid, config_path, RESOLVED_OUTPUT_DIR, state_was_recovered=True)
            
            # Mark that final outputs were generated
            _final_outputs_generated = True
        else:
            print("❌ No state available for output generation")
            
    finally:
        # Final safety net: always try to save current state if available
        if _current_state is not None and not _graceful_shutdown_initiated:
            try:
                print("\n🔒 Final safety check: ensuring state preservation...")
                # This serves as a backup in case the normal flow didn't save the state
                # Only save if we haven't already processed this state
                if not _final_outputs_generated:
                    generate_final_outputs(_current_state, _config_service, _maze_uuid, _config_path, _output_dir, state_was_recovered=True)
                    # Mark that we've generated final outputs
                    _final_outputs_generated = True
                else:
                    print("✅ Final outputs already generated for current state")
            except Exception as safety_error:
                print(f"⚠️  Final safety save failed: {safety_error}")
                # Last resort emergency save
                try:
                    _save_emergency_state(_current_state, _config_service, _maze_uuid, _config_path, _output_dir)
                except Exception as emergency_error:
                    print(f"❌ Emergency save also failed: {emergency_error}")


########################################################################################
# Helper Functions
########################################################################################

def generate_final_outputs(final_state, config_service, maze_uuid, config_path, RESOLVED_OUTPUT_DIR, state_was_recovered=False):
    """Generate final outputs including visualization and pickle files"""
    
    # Get final maze wrapper
    agent_id = list(final_state["maze_wrappers"].keys())[0]
    final_maze_wrapper = final_state["maze_wrappers"][agent_id]
    
    status_prefix = "RECOVERED" if state_was_recovered else "COMPLETED"
    
    if final_maze_wrapper.is_at_exit():
        print(f"SUCCESS! Agent found the exit!")
        print(f"Final position: {final_maze_wrapper.get_agent_position()}")
        print(f"Total moves: {final_maze_wrapper.get_move_count()}")
    else:
        print(f"Maze exploration ended")
        print(f"Final position: {final_maze_wrapper.get_agent_position()}")
        print(f"Total moves: {final_maze_wrapper.get_move_count()}")
    
    # Display and save final maze visualization with 50/50 layout using MazeAggregator
    try:
        print(f"\nDisplaying final multi-agent maze visualization ({status_prefix.lower()} state)...")
        
        # Create aggregator from final state
        from orchestrator_maze_implementation.visualization.maze_visualization import create_maze_aggregator_from_state, _MAZE_FIGURE_NUM
        final_aggregator = create_maze_aggregator_from_state(final_state)
        
        # Reuse the same figure window for final view
        plt.figure(_MAZE_FIGURE_NUM, figsize=(16, 8))
        plt.clf()  # Clear the figure
        
        # Create figure with 50/50 subplot layout for final view
        fig = plt.gcf()  # Get current figure
        ax_messages, ax_maze = fig.subplots(1, 2)
        
        # === LEFT SIDE: Final Multi-Agent Status Panel (50%) ===
        ax_messages.set_xlim(0, 1)
        ax_messages.set_ylim(0, 1)
        ax_messages.axis('off')
        ax_messages.set_title(f'Final Multi-Agent Status ({status_prefix})', fontsize=12, fontweight='bold')
        
        # Get final statistics
        final_agent_stats = final_aggregator.get_agent_statistics()
        final_global_stats = final_aggregator.get_global_statistics()
        final_agents_at_exit = final_aggregator.get_agents_at_exit()
        
        # Create status text
        status_lines = []
        status_lines.append(f"FINAL RESULTS ({status_prefix}):")
        status_lines.append(f"Total Agents: {final_global_stats['total_agents']}")
        status_lines.append(f"Agents at Exit: {final_global_stats['agents_at_exit']}")
        status_lines.append(f"Completion Rate: {final_global_stats['completion_rate']:.1f}%")
        status_lines.append(f"Total Moves: {final_global_stats['total_moves']}")
        status_lines.append(f"Success Rate: {final_global_stats['success_rate']:.1f}%")
        if state_was_recovered:
            status_lines.append(f"⚠️  STATE WAS RECOVERED FROM TIMEOUT/ERROR")
        status_lines.append("-" * 40)
        
        # Show individual agent final status
        for agent_id, stats in final_agent_stats.items():
            status_lines.append(f"Agent {agent_id}: {stats['moves']} moves, {'AT EXIT' if stats['at_exit'] else 'exploring'}")
        
        # Add recent messages if available
        if final_state.get('messages'):
            recent_messages = final_state['messages'][-3:] if len(final_state['messages']) > 3 else final_state['messages']
            status_lines.append("-" * 40)
            status_lines.append("Recent Messages:")
            for msg in recent_messages:
                content = str(msg.content)[:50] + "..." if len(str(msg.content)) > 50 else str(msg.content)
                status_lines.append(f"• {content}")
        
        # Display status with proper spacing
        full_text = '\n'.join(status_lines)
        ax_messages.text(0.05, 0.95, full_text, 
                       transform=ax_messages.transAxes,
                       fontsize=9,
                       verticalalignment='top',
                       horizontalalignment='left',
                       bbox=dict(facecolor='lightblue' if state_was_recovered else 'lightgreen', alpha=0.3, pad=10),
                       family='monospace')
        
        # === RIGHT SIDE: Final Multi-Agent Maze Panel (50%) ===
        try:
            # Draw final multi-agent maze directly
            height, width = final_aggregator.maze.shape
            
            # Draw base maze
            for y in range(height):
                for x in range(width):
                    cell = final_aggregator.maze[y, x]
                    color = final_aggregator.base_colors.get(cell, '#FFFFFF')
                    rect = matplotlib.patches.Rectangle((x, height-y-1), 1, 1, 
                                        facecolor=color, edgecolor='black', linewidth=0.5)
                    ax_maze.add_patch(rect)
            
            # Draw agent paths (complete move history) with new color scheme
            for i, (agent_id, agent) in enumerate(final_aggregator.agents.items()):
                # Draw visited paths as light pink squares (matching updated scheme)
                if len(agent.move_history) > 1:
                    # Draw visited path squares (excluding current position)
                    for j, (row, col) in enumerate(agent.move_history[:-1]):
                        visited_rect = matplotlib.patches.Rectangle((col, height - row - 1), 1, 1,
                                                    facecolor='#FFB6C1', edgecolor='none',
                                                    alpha=0.6)
                        ax_maze.add_patch(visited_rect)
                
                # Draw path markers
                color = final_aggregator.agent_colors[i % len(final_aggregator.agent_colors)]
                
                # Draw small path markers
                if len(agent.move_history) > 1:
                    # Calculate horizontal offset for this agent
                    offset_x = get_agent_offset_x(agent_id, i)
                    
                    # Draw path markers (excluding current position)
                    for j, (row, col) in enumerate(agent.move_history[:-1]):
                        circle = matplotlib.patches.Circle((col + 0.5 + offset_x, height - row - 0.5), 0.1, 
                                            facecolor=color, edgecolor=color, linewidth=1,
                                            alpha=0.8)
                        ax_maze.add_patch(circle)
            
            # Draw marked dead ends for all agents
            all_marked_dead_ends = set()
            for agent_id, agent in final_aggregator.agents.items():
                marked_dead_ends = agent.get_marked_dead_ends()
                for dead_row, dead_col in marked_dead_ends:
                    if (0 <= dead_row < height and 0 <= dead_col < width):
                        all_marked_dead_ends.add((dead_row, dead_col))
            
            # Draw dead end markings
            for dead_row, dead_col in all_marked_dead_ends:
                # Draw red square for dead end
                dead_rect = matplotlib.patches.Rectangle((dead_col, height - dead_row - 1), 1, 1,
                                        facecolor='#FF0000', edgecolor='#8B0000', linewidth=2,
                                        alpha=0.7)
                ax_maze.add_patch(dead_rect)
            
            # Draw starting positions
            for i, (agent_id, agent) in enumerate(final_aggregator.agents.items()):
                if agent.move_history:
                    start_row, start_col = agent.move_history[0]
                    color = final_aggregator.agent_colors[i % len(final_aggregator.agent_colors)]
                    
                    # Draw start marker as a square
                    square = matplotlib.patches.Rectangle((start_col + 0.2, height - start_row - 0.8), 0.6, 0.6,
                                            facecolor='white', edgecolor=color, linewidth=3)
                    ax_maze.add_patch(square)
                    
                    # Add 'S' label
                    ax_maze.text(start_col + 0.5, height - start_row - 0.5, 'S', 
                                ha='center', va='center', fontsize=8, fontweight='bold', color=color)
            
            # Draw final agent positions
            for i, (agent_id, agent) in enumerate(final_aggregator.agents.items()):
                row, col = agent.get_agent_position()
                color = final_aggregator.agent_colors[i % len(final_aggregator.agent_colors)]
                
                # Calculate horizontal offset for this agent
                offset_x = get_agent_offset_x(agent_id, i)
                
                # Special styling for agents at exit
                if agent_id in final_agents_at_exit:
                    circle = matplotlib.patches.Circle((col + 0.5 + offset_x, height - row - 0.5), 0.35, 
                                        facecolor=color, edgecolor='gold', linewidth=4)
                else:
                    circle = matplotlib.patches.Circle((col + 0.5 + offset_x, height - row - 0.5), 0.3, 
                                        facecolor=color, edgecolor='black', linewidth=2)
                ax_maze.add_patch(circle)
                
                # Add agent label
                label = agent_id.replace('agent_', 'A') if 'agent_' in agent_id else agent_id[:2]
                ax_maze.text(col + 0.5 + offset_x, height - row - 0.5, label, 
                            ha='center', va='center', fontsize=8, fontweight='bold', color='black')
            
            # Set axis properties
            ax_maze.set_xlim(0, width)
            ax_maze.set_ylim(0, height)
            ax_maze.set_aspect('equal')
            
            # Set title
            completion_rate = len(final_agents_at_exit) / len(final_aggregator.agents) * 100
            ax_maze.set_title(f'Final Multi-Agent Maze ({completion_rate:.0f}% Complete)', 
                                fontsize=12, fontweight='bold')
            
            # Remove ticks for cleaner look
            ax_maze.set_xticks([])
            ax_maze.set_yticks([])
            
            # Add legend with agent colors and final stats
            legend_elements = [
                matplotlib.patches.Rectangle((0,0),1,1, facecolor='#000000', label='Wall'),
                matplotlib.patches.Rectangle((0,0),1,1, facecolor='#FFFFFF', label='Open Path'),
                matplotlib.patches.Rectangle((0,0),1,1, facecolor='#90EE90', label='Exit'),
                matplotlib.patches.Rectangle((0,0),1,1, facecolor='#808080', label='Frame'),
            ]
            
            # Add agents to legend with their final stats
            for i, agent_id in enumerate(final_aggregator.get_agent_ids()):
                agent = final_aggregator.get_agent(agent_id)
                color = final_aggregator.agent_colors[i % len(final_aggregator.agent_colors)]
                status = " (EXIT)" if agent_id in final_agents_at_exit else ""
                moves = f" [{len(agent.move_history)-1}m]" if len(agent.move_history) > 1 else ""
                legend_elements.append(
                    plt.Line2D([0], [0], marker='o', color='w', markerfacecolor=color, 
                                markersize=8, label=f'{agent_id}{status}{moves}')
                )
            
            ax_maze.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(1.3, 1), fontsize=8)
            
        except Exception as e:
            # Fallback: display error message
            ax_maze.text(0.5, 0.5, f'Final multi-agent visualization error:\n{str(e)}', 
                        transform=ax_maze.transAxes,
                        fontsize=10, ha='center', va='center',
                        bbox=dict(facecolor='red', alpha=0.3))
            ax_maze.set_title('Final Multi-Agent Visualization Error', fontsize=12, fontweight='bold')
        
        # Adjust layout, save, and then try to display
        plt.tight_layout()

        # Generate timestamp for unique filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # Save multi-agent visualization to file regardless of display availability
        maze_viz_filepath = os.path.join(RESOLVED_OUTPUT_DIR, f"final_multi_agent_maze_state{'_recovered' if state_was_recovered else ''}_{timestamp}.png")
        final_aggregator.save_combined_visualization(maze_viz_filepath)
        print(f"Final multi-agent maze state saved as: {maze_viz_filepath}")

        # Try to display non-blocking if a GUI backend is available (gated by env flag)
        if _get_safe_mode_from_config_or_env(config_service):
            try:
                plt.show(block=False)
                plt.pause(5)
                plt.draw()
            except Exception as display_error:
                print(f"⚠️  Could not display final visualization (GUI backend unavailable?): {display_error}")
        else:
            # Original behavior
            plt.show(block=False)
            plt.pause(5)
            plt.draw()
    except Exception as e:
        print(f"Could not display/save final maze visualization: {e}")
    
    # Display conversation summary
    try:
        display_conversation_summary(final_state)
    except Exception as e:
        print(f"Could not display conversation summary: {e}")
    
    # Show full conversation history automatically at the end
    try:
        print(f"\nFULL CONVERSATION HISTORY ({status_prefix})")
        print("=" * 60)
        
        for i, msg in enumerate(final_state.get("messages", [])):
            if isinstance(msg, (SystemMessage, HumanMessage, AIMessage, ToolMessage)):
                role = msg.__class__.__name__
                content = msg.content
                print(f"\n{i+1}. [{role}]: {content}")
        
        print("=" * 60)
    except Exception as e:
        print(f"Could not display conversation history: {e}")
    
    # Save final state and config to pickle file with maze UUID and timestamp
    try:
        # Create timestamp string for filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # Create filename with maze UUID and timestamp
        pickle_filename = f"maze_results_{maze_uuid}_{timestamp}{'_recovered' if state_was_recovered else ''}.pickle"
        
        # Create outputs directory (already created earlier) and save there
        pickle_filepath = os.path.join(RESOLVED_OUTPUT_DIR, pickle_filename)
        
        # Prepare data to save
        save_data = {
            "final_state": final_state,
            "config": config_service.get_config(),
            "maze_uuid": maze_uuid,
            "timestamp": timestamp,
            "recovered_from_timeout": state_was_recovered,
            "execution_metadata": {
                "config_file_path": os.path.abspath(config_path),
                "total_agents": len(final_state.get("maze_wrappers", {})),
                "recursion_limit": config_service.get_recursion_limit(),
                "steps_per_agent": config_service.get_steps_per_agent()
            }
        }
        
        # Save to pickle file
        with open(pickle_filepath, 'wb') as f:
            pickle.dump(save_data, f, protocol=pickle.HIGHEST_PROTOCOL)
        
        print(f"\n✅ Results saved to pickle file: {pickle_filepath}")
        print(f"   - Final state with {len(final_state.get('messages', []))} messages")
        print(f"   - Amount of turns: {final_state.get('turn_count', 'unknown')} of {config_service.get_max_total_steps()}")
        if state_was_recovered:
            print(f"   - ⚠️  State was recovered from timeout/error")
        print(f"   - Configuration settings")
        print(f"   - Execution metadata")
        
    except Exception as e:
        print(f"❌ Error saving results to pickle file: {e}")


def optimize_exit_placement(maze_wrapper):
    """
    Reposition the exit to be far from the starting position, preferably at the end of a path or in a corner.
    
    Args:
        maze_wrapper: MazeWrapper instance with the maze to optimize
    """
    from collections import deque
    
    maze = maze_wrapper.maze
    start_pos = maze_wrapper.get_agent_position()
    height, width = maze.shape
    
    # Find all open cells (including current exit)
    open_cells = []
    current_exit = None
    for r in range(height):
        for c in range(width):
            if maze[r, c] in ['O', 'E']:
                if maze[r, c] == 'E':
                    current_exit = (r, c)
                    # Convert current exit back to open cell for analysis
                    maze[r, c] = 'O'
                open_cells.append((r, c))
    
    if not open_cells:
        print("⚠️  No open cells found for exit placement")
        return
    
    # Calculate Manhattan distance from start to all open cells
    def manhattan_distance(pos1, pos2):
        return abs(pos1[0] - pos2[0]) + abs(pos1[1] - pos2[1])
    
    # Calculate actual path distance using BFS
    def calculate_path_distance(start, target):
        if start == target:
            return 0
            
        queue = deque([(start, 0)])
        visited = {start}
        
        directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]  # right, left, down, up
        
        while queue:
            (r, c), dist = queue.popleft()
            
            for dr, dc in directions:
                nr, nc = r + dr, c + dc
                
                if (0 <= nr < height and 0 <= nc < width and 
                    (nr, nc) not in visited and maze[nr, nc] == 'O'):
                    
                    if (nr, nc) == target:
                        return dist + 1
                    
                    visited.add((nr, nc))
                    queue.append(((nr, nc), dist + 1))
        
        return float('inf')  # No path found
    
    # Score potential exit positions
    def score_exit_position(pos):
        r, c = pos
        score = 0
        
        # 1. Distance from start (higher is better)
        path_dist = calculate_path_distance(start_pos, pos)
        if path_dist == float('inf'):
            return -1000  # Unreachable positions get very low score
        
        manhattan_dist = manhattan_distance(start_pos, pos)
        score += path_dist * 10 + manhattan_dist * 5
        
        # 2. Prefer positions near corners or edges
        edge_proximity = 0
        if r <= 2 or r >= height - 3:  # Near top or bottom edge
            edge_proximity += 15
        if c <= 2 or c >= width - 3:   # Near left or right edge  
            edge_proximity += 15
        if (r <= 2 or r >= height - 3) and (c <= 2 or c >= width - 3):  # Corner bonus
            edge_proximity += 25
        score += edge_proximity
        
        # 3. Prefer dead ends or positions with few connections
        directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]
        connection_count = 0
        for dr, dc in directions:
            nr, nc = r + dr, c + dc
            if (0 <= nr < height and 0 <= nc < width and maze[nr, nc] == 'O'):
                connection_count += 1
        
        # Dead ends (1 connection) get highest bonus, then 2 connections
        if connection_count == 1:
            score += 30  # Dead end bonus
        elif connection_count == 2:
            score += 15  # Path end bonus
        elif connection_count >= 4:
            score -= 10  # Junction penalty
            
        # 4. Avoid positions too close to the center
        center_r, center_c = height // 2, width // 2
        center_dist = manhattan_distance(pos, (center_r, center_c))
        score += center_dist * 2
        
        return score
    
    # Evaluate all potential exit positions
    best_score = -float('inf')
    best_positions = []
    
    print(f"🔍 Evaluating {len(open_cells)} potential exit positions...")
    print(f"   Starting position: {start_pos}")
    
    for pos in open_cells:
        if pos == start_pos:  # Don't place exit at start
            continue
            
        score = score_exit_position(pos)
        if score > best_score:
            best_score = score
            best_positions = [pos]
        elif score == best_score and score > 0:
            best_positions.append(pos)
    
    if not best_positions:
        print("⚠️  No suitable exit positions found, keeping original")
        if current_exit:
            maze[current_exit[0], current_exit[1]] = 'E'
        return
    
    # Choose the best position (prefer one that's furthest from start if tied)
    if len(best_positions) > 1:
        best_pos = max(best_positions, key=lambda pos: manhattan_distance(start_pos, pos))
    else:
        best_pos = best_positions[0]
    
    # Place the new exit
    maze[best_pos[0], best_pos[1]] = 'E'
    
    # Calculate final metrics
    final_path_dist = calculate_path_distance(start_pos, best_pos)
    final_manhattan_dist = manhattan_distance(start_pos, best_pos)
    
    print(f"✅ Exit repositioned to: {best_pos}")
    print(f"   Path distance from start: {final_path_dist}")
    print(f"   Manhattan distance: {final_manhattan_dist}")
    print(f"   Position score: {best_score}")
    
    # Check if it's a dead end
    directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]
    connections = sum(1 for dr, dc in directions 
                     if (0 <= best_pos[0] + dr < height and 0 <= best_pos[1] + dc < width and 
                         maze[best_pos[0] + dr, best_pos[1] + dc] == 'O'))
    if connections == 1:
        print("   🎯 Exit placed at dead end!")
    elif connections == 2:
        print("   🎯 Exit placed at path end!")


def initialize_state(num_agents: int, maze_uuid: str) -> MazeState:
    """
    Initialize MazeState with all fields set to safe defaults, including per-agent messages.
    
    Args:
        num_agents: Number of agents to create
        maze_uuid: UUID of the maze to load for all agents
        
    Returns:
        MazeState: Initialized state with the specified number of agents and maze
    """
    maze_wrappers = {}
    agent_messages = {}
    agent_positions = {} 
    previously_visited_tiles = {}
    all_agents = []
    known_openings = {}
    free_energy_metrics = {}
    entropy_history = {}
    agent_backtracking_state = {}
    
    for i in range(num_agents):
        agent_id = f"agent_{i}"
        maze_wrapper = create_maze_by_uuid(maze_uuid)  # Use configurable maze UUID
        
        # Optimize exit placement to be far from start and preferably at path ends
        print(f"🎯 Optimizing exit placement for {agent_id}...")
        optimize_exit_placement(maze_wrapper)
        
        maze_wrappers[agent_id] = maze_wrapper
        agent_messages[agent_id] = []  # Initialize empty list for each agent's messages
        agent_positions[agent_id] = maze_wrapper.get_agent_position()
        previously_visited_tiles[agent_id] = maze_wrapper.move_history.copy()
        all_agents.append(agent_id)
        known_openings[agent_id] = []
        free_energy_metrics[agent_id] = []
        entropy_history[agent_id] = []
        agent_backtracking_state[agent_id] = {}
        
    # Initialize current agent using the rotation function
    initial_current_agent = get_current_agent_by_turn(0, all_agents)

    return MazeState(
        maze_wrappers=maze_wrappers,
        messages=[],  # Assuming this is a global field; adjust if not
        agent_messages=agent_messages,  # Now populated with keys for each agent
        system_task="",
        plan=[],
        plan_completed=False,
        step_index=0,
        maze_exit_found=False,
        exit_position=[],
        winning_agent="",
        turn_count=0,

        # Agent-specific fields
        current_agent=initial_current_agent,  # Use rotation function
        all_agents=all_agents,
        known_openings=known_openings,
        # agent_positions=agent_positions,
        # previously_visited_tiles=previously_visited_tiles,
        
        # Backtracking state management
        agent_backtracking_state={},
        
        # Initialize orchestrator fields
        orchestrator_guidance=None,
        shared_knowledge=None,
        turn_complete= False,
        free_energy_metrics={},
        entropy_history=[],
    )

########################################################################################
# Call Main
########################################################################################

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Multi-Agent Maze Runner",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    
    parser.add_argument(
        "--config", "-c",
        type=str,
        default=None,
        help="Path to configuration file (default: use built-in default)"
    )
    
    parser.add_argument(
        "--output_dir", "-o",
        type=str,
        default=None,
        help="Optional directory to save final artifacts (images, pickles). Defaults to '<config_dir>/outputs'"
    )
    
    args = parser.parse_args()
    main(config_path=args.config, output_dir=args.output_dir)