import time
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.colors import ListedColormap
from matplotlib.font_manager import FontProperties
import numpy as np
import os
from typing import List, Tuple, Optional, Dict, Any
from orchestrator_maze_implementation.state.maze_state import MazeState
from orchestrator_maze_implementation.utils.maze_utils.maze_aggregator import MazeAggregator
from orchestrator_maze_implementation.utils.maze_utils.maze_generation import MazeGenerator
from orchestrator_maze_implementation.config.config_service import config
from orchestrator_maze_implementation.visualization.fe_metrics_graph import create_fe_metrics_graph

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

def get_agent_offset_x(agent_id: str, agent_index: int) -> float:
    """
    Calculate horizontal offset for agent visualization based on agent ID.
    
    Args:
        agent_id: Unique identifier for the agent (e.g., 'agent_0', 'agent_1')
        agent_index: Index of the agent in the iteration (0, 1, 2, ...)
    
    Returns:
        Horizontal offset in pixels (converted to matplotlib coordinates)
    """
    # Extract numeric ID from agent_id if possible, otherwise use agent_index
    try:
        if 'agent_' in agent_id:
            numeric_id = int(agent_id.replace('agent_', ''))
        else:
            # Try to extract number from end of string
            import re
            match = re.search(r'(\d+)$', agent_id)
            numeric_id = int(match.group(1)) if match else agent_index
    except (ValueError, AttributeError):
        numeric_id = agent_index
    
    # Create offset pattern: alternate between positive and negative offsets
    # Agent 0: 0, Agent 1: +10px, Agent 2: -10px, Agent 3: +20px, Agent 4: -20px, etc.
    if numeric_id == 0:
        offset_pixels = 0
    elif numeric_id % 2 == 1:  # Odd agents get positive offset
        offset_pixels = ((numeric_id + 1) // 2) * 10
    else:  # Even agents get negative offset
        offset_pixels = -(numeric_id // 2) * 10
    
    # Convert pixels to matplotlib coordinate units
    # Assume each cell is approximately 50 pixels wide for conversion
    return offset_pixels / 50.0

def configure_emoji_font():
    """Configure matplotlib to use a font that supports emojis with fallback strategy."""
    try:
        import matplotlib.font_manager as fm
        
        # Priority list of emoji-supporting fonts (macOS, Windows, Linux)
        emoji_fonts = [
            'Apple Color Emoji',    # macOS primary
            'SF Pro Text',          # macOS fallback
            'Helvetica Neue',       # macOS fallback
            'Arial Unicode MS',     # Cross-platform
            'Segoe UI Emoji',       # Windows primary
            'Segoe UI',             # Windows fallback
            'Noto Color Emoji',     # Linux primary
            'Noto Sans',            # Linux fallback
            'Liberation Sans',      # Linux alternative
            'Arial',                # Universal fallback
        ]
        
        # Get available fonts on system
        available_fonts = [f.name for f in fm.fontManager.ttflist]
        
        # Find and set the best available emoji font
        for font_name in emoji_fonts:
            if font_name in available_fonts:
                plt.rcParams['font.family'] = font_name
                print(f"✅ Emoji font configured: {font_name}")
                return font_name
        
        # Ultimate fallback - disable font warnings and use text alternatives
        import warnings
        warnings.filterwarnings('ignore', category=UserWarning, module='matplotlib')
        print("⚠️  No optimal emoji font found - using text alternatives")
        return None
        
    except Exception as e:
        print(f"⚠️  Font configuration error: {e} - using system default")
        return None

# Configure emoji support when module is imported
emoji_font = configure_emoji_font()

# Global figure number for reusing the same window
_MAZE_FIGURE_NUM = 1

def create_maze_aggregator_from_state(state: MazeState) -> MazeAggregator:
    """
    Create a MazeAggregator from the current MazeState.
    
    Args:
        state: Current MazeState with multiple agents
        
    Returns:
        MazeAggregator instance with all agents loaded
    """
    # Get the first agent to extract the maze generator
    first_agent_id = list(state["maze_wrappers"].keys())[0]
    first_maze_wrapper = state["maze_wrappers"][first_agent_id]
    
    # Create aggregator with the same maze generator
    aggregator = MazeAggregator(first_maze_wrapper.maze_generator)
    
    # Add all agents to the aggregator
    for agent_id, maze_wrapper in state["maze_wrappers"].items():
        aggregator.add_agent(agent_id, maze_wrapper=maze_wrapper, avoid_occupied=False)
    
    return aggregator

def display_maze_state(state: MazeState, enable_visualization=None):
    """Display the current maze state using matplotlib visualization with multi-agent support"""
    
    # Use config value if not provided
    if enable_visualization is None:
        try:
            enable_visualization = config.get_enable_maze_visualization()
        except RuntimeError:
            # Config not loaded, use default
            print("⚠️  Configuration not loaded, using default visualization setting (enabled)")
            enable_visualization = True
    
    MAZE_VIS_SAFE_MODE = _get_safe_mode()
    if MAZE_VIS_SAFE_MODE:
        # Safe mode: don't force a specific GUI backend; control draw/show manually
        plt.ioff()
    else:
        # Original behavior: prefer TkAgg if possible, and hide root window to avoid focus stealing
        try:
            matplotlib.use('TkAgg')
        except Exception:
            # Fallback: if Tk isn't available, try Qt; otherwise keep default
            try:
                matplotlib.use('Qt5Agg')
            except Exception:
                pass
        plt.ioff()
        try:
            import tkinter as tk
            tk.Tk().withdraw()
        except Exception:
            pass

    print("\n" + "="*50)
    print("CURRENT MAZE STATE")
    print("="*50)
    
    # Create aggregator from state
    aggregator = create_maze_aggregator_from_state(state)
    
    # Display basic info for all agents
    current_agent = state.get("current_agent", "unknown")
    print(f"Current active agent: {current_agent}")
    print(f"Total agents: {len(state['maze_wrappers'])}")
    print(f"Messages: {len(state['messages'])}")
    
    # Show individual agent status
    agent_stats = aggregator.get_agent_statistics()
    agents_at_exit = aggregator.get_agents_at_exit()
    
    for agent_id, stats in agent_stats.items():
        status = "🎉 AT EXIT" if agent_id in agents_at_exit else "🏃 NAVIGATING"
        current_marker = " (CURRENT)" if agent_id == current_agent else ""
        print(f"Agent {agent_id}{current_marker}: {stats['position']} - {status}")
        print(f"  Moves: {stats['moves']}, Failed: {stats['failed_moves']}")
    
    # Show global statistics
    global_stats = aggregator.get_global_statistics()
    print(f"Global: {global_stats['total_moves']} moves, {global_stats['agents_at_exit']}/{global_stats['total_agents']} at exit")
    
    print("="*50)
    
    if not enable_visualization:
        print("Visualization disabled - showing text display only. To turn on visualization set enable_maze_visualization to True in config.yaml")
        # Show text display for current agent
        current_agent_wrapper = state["maze_wrappers"][current_agent]
        maze_display = current_agent_wrapper.get_maze_with_agent()
        time.sleep(0.1)
        return
    
    # Debug info for visualization
    print(f"🎨 Matplotlib backend: {matplotlib.get_backend()}")
    print(f"🎨 Enable visualization: {enable_visualization}")
    print(f"🎨 Safe mode: {MAZE_VIS_SAFE_MODE}")
    print(f"🎨 Display environment: {'DISPLAY' in os.environ}")
    
    # Create and display matplotlib visualization with 50/50 layout
    try:
        print("Generating multi-agent maze visualization...")
        
        # Clear any existing figure content and reuse the same window
        fig = plt.figure(_MAZE_FIGURE_NUM, figsize=(20, 12))
        plt.clf()  # Clear the figure
        
        # Set elegant background color
        fig.patch.set_facecolor('#FAFBFC')  
        
        # Create custom gridspec for asymmetric layout with space for legend
        gs = fig.add_gridspec(2, 2, width_ratios=[1, 2], height_ratios=[1, 1], 
                             hspace=0.15, wspace=0.18,  # Increased spacing between left and right panels
                             left=0.02, right=0.82, top=0.95, bottom=0.05)  # Leave space for legend on right
        
        # Top left: Messages Panel (25%)
        ax_messages = fig.add_subplot(gs[0, 0])
        
        # Bottom left: Free Energy Graph (25%)  
        ax_fe_graph = fig.add_subplot(gs[1, 0])
        
        # Right side: Maze Panel (50% - spans both rows)
        ax_maze = fig.add_subplot(gs[:, 1])
        
        # === LEFT SIDE: Messages Panel (50%) ===
        ax_messages.set_xlim(0, 1)
        ax_messages.set_ylim(0, 1)
        ax_messages.axis('off')
        ax_messages.set_title('Multi-Agent Status & Messages', fontsize=12, fontweight='bold')
        
        # Create content for messages panel
        msg_lines = []
        
        # Show current agent info and rotation
        all_agents = state.get('all_agents', list(state['maze_wrappers'].keys()))
        current_idx = all_agents.index(current_agent) if current_agent in all_agents else 0
        next_agent = all_agents[(current_idx + 1) % len(all_agents)]
        
        msg_lines.append(f"[ACTIVE] Agent: {current_agent}")
        msg_lines.append(f"[NEXT] Agent: {next_agent}")
        msg_lines.append(f"[TURN] Count: {state.get('turn_count', 0)}")
        msg_lines.append("-" * 40)
        
        # Show agent statuses
        for agent_id, stats in agent_stats.items():
            marker = ">>> " if agent_id == current_agent else "    "
            status = "EXIT" if agent_id in agents_at_exit else "ACTIVE"
            msg_lines.append(f"{marker}{agent_id}: {stats['position']} ({status})")
        
        msg_lines.append("-" * 40)
        
        # Show recent messages if available
        if state.get('agent_messages') and current_agent in state['agent_messages']:
            agent_messages = state["agent_messages"][current_agent]
            
            # Filter for useful messages
            useful_msgs = []
            for msg in agent_messages:
                msg_type = msg.__class__.__name__.replace('Message', '')
                
                if msg_type == 'AI' and (msg.content or (hasattr(msg, 'tool_calls') and msg.tool_calls)):
                    useful_msgs.append(msg)
                elif msg_type == 'Tool' and msg.content:
                    if any(keyword in msg.content.lower() for keyword in 
                           ['position', 'moved', 'view', 'possible', 'dead end', 'exit', 'error']):
                        useful_msgs.append(msg)
                elif msg_type in ['System', 'Human'] and msg.content:
                    useful_msgs.append(msg)
            
            # Show last few messages
            try:
                message_limit = config.get_message_history_limit()
            except RuntimeError:
                message_limit = 4  # Default value if config not loaded
            latest_msgs = useful_msgs[-message_limit:]
            
            msg_lines.append(f"Recent messages ({current_agent}):")
            for i, msg in enumerate(latest_msgs):
                role = msg.__class__.__name__.replace('Message', '')
                content = msg.content
                if len(content) > 60:
                    content = content[:57] + "..."
                msg_lines.append(f"  {i+1}. [{role}]: {content}")
        
        # Display messages with elegant styling
        full_text = '\n'.join(msg_lines)
        ax_messages.text(0.05, 0.95, full_text, 
                       transform=ax_messages.transAxes,
                       fontsize=9,
                       verticalalignment='top',
                       horizontalalignment='left',
                       bbox=dict(facecolor='#F8F9FA', alpha=0.9, pad=12, 
                               edgecolor='#DEE2E6', linewidth=1,
                               boxstyle='round,pad=0.5'),
                       family='monospace',
                       color='#495057')
        
        # Add elegant panel styling
        ax_messages.set_facecolor('#FAFBFC')
        ax_messages.spines['top'].set_visible(False)
        ax_messages.spines['right'].set_visible(False)
        ax_messages.spines['bottom'].set_visible(False)
        ax_messages.spines['left'].set_visible(False)
        
        # === BOTTOM LEFT: Free Energy Metrics Graph (25%) ===
        try:
            # Get entropy history from state (might be in different locations)
            entropy_history = []
            if 'entropy_history' in state:
                entropy_history = state['entropy_history']
            
            # Ensure entropy_history is a list and handle None case
            if entropy_history is None:
                entropy_history = []
            
            # Get agent colors for consistency with maze
            aggregator = create_maze_aggregator_from_state(state)
            agent_colors = aggregator.agent_colors
            
            # Create the FE metrics graph
            create_fe_metrics_graph(
                ax=ax_fe_graph,
                entropy_history=entropy_history,
                current_agent=current_agent,
                all_agents=all_agents,
                turn_count=state.get('turn_count', 0),
                agent_colors=agent_colors
            )
            
        except Exception as fe_error:
            # Fallback: display placeholder message
            ax_fe_graph.text(0.5, 0.5, f'Free Energy Graph Error:\n{str(fe_error)}\n\nGraph will show once\nbenchmarking data is available', 
                           transform=ax_fe_graph.transAxes,
                           fontsize=10, ha='center', va='center',
                           bbox=dict(facecolor='lightyellow', alpha=0.7))
            ax_fe_graph.set_title('Free Energy Metrics', fontsize=11, fontweight='bold')
        
        # === RIGHT SIDE: Multi-Agent Maze Panel (50%) ===
        try:
            # Draw multi-agent maze directly without creating intermediate figure
            height, width = aggregator.maze.shape
            
            # Draw base maze
            for y in range(height):
                for x in range(width):
                    cell = aggregator.maze[y, x]
                    color = aggregator.base_colors.get(cell, '#FFFFFF')
                    rect = plt.Rectangle((x, height-y-1), 1, 1, 
                                       facecolor=color, edgecolor='black', linewidth=0.5)
                    ax_maze.add_patch(rect)
            
            # Draw agent paths (move history) with light pink visited path color
            for i, (agent_id, agent) in enumerate(aggregator.agents.items()):
                # Draw visited paths as light pink squares (matching image)
                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 = plt.Rectangle((col, height - row - 1), 1, 1,
                                                   facecolor='#FFB6C1', edgecolor='none',
                                                   alpha=0.6)
                        ax_maze.add_patch(visited_rect)
                
                # Draw path markers as small circles for clarity
                color = aggregator.agent_colors[i % len(aggregator.agent_colors)]
                path_color = aggregator._lighten_color(color, 0.6)  # Lighter version for path
                
                if len(agent.move_history) > 1:
                    # Calculate horizontal offset for this agent
                    offset_x = get_agent_offset_x(agent_id, i)
                    
                    # Draw small path markers (excluding current position)
                    for j, (row, col) in enumerate(agent.move_history[:-1]):
                        circle = plt.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 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 = plt.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(aggregator.agents.items()):
                if agent.move_history:
                    start_row, start_col = agent.move_history[0]
                    color = aggregator.agent_colors[i % len(aggregator.agent_colors)]
                    
                    # Draw start marker as a square
                    square = plt.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 current agent positions (on top)
            for i, (agent_id, agent) in enumerate(aggregator.agents.items()):
                row, col = agent.get_agent_position()
                color = aggregator.agent_colors[i % len(aggregator.agent_colors)]
                
                # Calculate horizontal offset for this agent
                offset_x = get_agent_offset_x(agent_id, i)
                
                # Highlight current active agent with larger circle and white border
                if agent_id == current_agent:
                    circle = plt.Circle((col + 0.5 + offset_x, height - row - 0.5), 0.35, 
                                      facecolor=color, edgecolor='white', linewidth=3)
                else:
                    circle = plt.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 with better contrast
                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 with current agent highlight
            ax_maze.set_title(f'Multi-Agent Maze (Active: {current_agent})', 
                             fontsize=12, fontweight='bold')
            
            # Remove ticks for cleaner look
            ax_maze.set_xticks([])
            ax_maze.set_yticks([])
            
            # Add legend
            legend_elements = [
                plt.Rectangle((0,0),1,1, facecolor='#000000', label='Wall'),
                plt.Rectangle((0,0),1,1, facecolor='#FFFFFF', label='Open Path'),
                plt.Rectangle((0,0),1,1, facecolor='#FFD700', label='Exit'),
                plt.Rectangle((0,0),1,1, facecolor='#808080', label='Frame'),
                plt.Rectangle((0,0),1,1, facecolor='#FFB6C1', label='Visited Path'),
                plt.Rectangle((0,0),1,1, facecolor='#FF0000', edgecolor='#8B0000', linewidth=2, label='Dead End'),
            ]
            
            # Add agents to legend with their colors and path info
            for i, agent_id in enumerate(aggregator.get_agent_ids()):
                agent = aggregator.get_agent(agent_id)
                color = aggregator.agent_colors[i % len(aggregator.agent_colors)]
                status = " (EXIT)" if agent_id in agents_at_exit else ""
                active = " (ACTIVE)" if agent_id == current_agent else ""
                path_info = 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}{active}{status}{path_info}')
                )
            
            ax_maze.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(1.25, 1), fontsize=8)
            
            # Add lock icon if any agent is in backtracking lock mode
            agents_in_lock_mode = []
            backtracking_state = state.get("agent_backtracking_state", {}) or {}
            for agent_id in aggregator.get_agent_ids():
                agent_backtrack = backtracking_state.get(agent_id, {})
                if agent_backtrack.get("lock_mode", False):
                    agents_in_lock_mode.append(agent_id)
            
            if agents_in_lock_mode:
                # Display lock icon underneath the legend, aligned with upper right legend position
                lock_text = f"🔒 Backtrack Lock: {', '.join(agents_in_lock_mode)}"
                ax_maze.text(1.25, 0.85, lock_text, transform=ax_maze.transAxes, 
                           fontsize=9, ha='right', va='top',
                           bbox=dict(boxstyle="round,pad=0.3", facecolor='#FFE4B5', alpha=0.8, edgecolor='#DAA520'))
            
        except Exception as e:
            # Fallback: display error message
            ax_maze.text(0.5, 0.5, f'Multi-agent maze 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('Multi-Agent Maze Visualization Error', fontsize=12, fontweight='bold')
            print(f"Maze visualization error: {e}")
        
        # Adjust layout and display
        plt.tight_layout()
        
        # Try to display the plot
        try:
            plt.show(block=False)
            plt.pause(0.5)  # 0.5 second delay as requested
            plt.draw()  # Force redraw
            print("✅ Visualization displayed successfully")
        except Exception as display_error:
            print(f"⚠️  Display error: {display_error}")
            # Still save the plot even if display fails
            plt.savefig('current_maze_state.png', dpi=150, bbox_inches='tight')
            print("📁 Saved visualization as 'current_maze_state.png'")
        
        print("Multi-agent maze visualization displayed!")
        
    except Exception as e:
        print(f"Could not display multi-agent maze visualization: {e}")
        print("Falling back to text display...")
        
        # Fallback to text display for current agent
        current_agent_wrapper = state["maze_wrappers"][current_agent]
        maze_display = current_agent_wrapper.get_maze_with_agent()
        for row in maze_display:
            print(''.join(row))
        
        print("LEGEND: A=Agent, E=Exit, O=Open path, W=Wall, X=Boundary")
        time.sleep(0.5)  # Still apply 0.5 second delay for text display

