"""
Display utilities for colored agent output in the terminal.

Provides automatic color assignment to agents for better visualization
of multi-agent conversations.
"""

from typing import Dict, Optional
from synqed.message_model import StructuredMessage


# ANSI color codes for terminal output
class Colors:
    """ANSI color codes for terminal output."""
    # Bright, distinct colors for agents
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    MAGENTA = '\033[95m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    
    # Additional colors
    ORANGE = '\033[38;5;208m'
    PURPLE = '\033[38;5;141m'
    PINK = '\033[38;5;213m'
    LIME = '\033[38;5;154m'
    TEAL = '\033[38;5;51m'
    
    # System colors
    GRAY = '\033[90m'
    BOLD = '\033[1m'
    
    # Reset
    RESET = '\033[0m'
    END = '\033[0m'


# Predefined color palette for agents
AGENT_COLORS = [
    Colors.RED,
    Colors.GREEN,
    Colors.CYAN,
    Colors.MAGENTA,
    Colors.YELLOW,
    Colors.BLUE,
    Colors.ORANGE,
    Colors.PURPLE,
    Colors.PINK,
    Colors.LIME,
    Colors.TEAL,
]


class ColoredAgentDisplay:
    """
    Manages colored display output for agents in a workspace.
    
    Automatically assigns distinct colors to each agent and provides
    formatted output for messages.
    
    Example:
        ```python
        display = ColoredAgentDisplay()
        workspace = synqed.create_workspace(
            name="My Workspace",
            task="Solve a problem",
            agents=[...],
            on_message=display.print_message,  # Use the colored printer
        )
        ```
    """
    
    def __init__(self, enabled: bool = True):
        """
        Initialize the colored display.
        
        Args:
            enabled: Whether to use colors (disable for non-terminal output)
        """
        self.enabled = enabled
        self.agent_colors: Dict[str, str] = {}
        self.color_index = 0
    
    def _assign_color(self, agent_id: str) -> str:
        """Assign a color to an agent."""
        if agent_id not in self.agent_colors:
            color = AGENT_COLORS[self.color_index % len(AGENT_COLORS)]
            self.agent_colors[agent_id] = color
            self.color_index += 1
        return self.agent_colors[agent_id]
    
    def get_agent_color(self, agent_id: str) -> str:
        """
        Get the color code for an agent.
        
        Args:
            agent_id: The agent's ID
            
        Returns:
            ANSI color code string
        """
        if not self.enabled:
            return ""
        return self._assign_color(agent_id)
    
    def print_message(self, message: StructuredMessage) -> None:
        """
        Print a message with colored agent names.
        
        This method can be used as an on_message callback for workspaces.
        
        Args:
            message: The structured message to print
        """
        sender_id = message.sender_id
        sender_name = message.sender_name
        content = message.content.strip()
        
        # Skip system initialization/phase messages for cleaner output
        if sender_id == "system":
            # Only show important system messages
            if any(keyword in content for keyword in ["TERMINATED", "FAILED", "ERROR"]):
                print(f"{Colors.GRAY}[System] {content}{Colors.RESET}")
            return
        
        # Get agent color
        if self.enabled:
            color = self._assign_color(sender_id)
            print(f"{color}[{sender_name}]{Colors.RESET} {content}")
        else:
            print(f"[{sender_name}] {content}")
    
    def print_message_verbose(self, message: StructuredMessage) -> None:
        """
        Print a message with full details (type, phase, etc.).
        
        Args:
            message: The structured message to print
        """
        sender_id = message.sender_id
        sender_name = message.sender_name
        content = message.content.strip()
        msg_type = message.message_type.value
        phase = message.metadata.phase
        
        if sender_id == "system":
            print(f"{Colors.GRAY}[System | {phase}] {content}{Colors.RESET}")
        else:
            if self.enabled:
                color = self._assign_color(sender_id)
                print(f"{color}[{sender_name}]{Colors.RESET} {Colors.GRAY}({msg_type} | {phase}){Colors.RESET}")
                print(f"  {content}")
            else:
                print(f"[{sender_name}] ({msg_type} | {phase})")
                print(f"  {content}")
    
    def format_agent_name(self, agent_id: str, agent_name: str) -> str:
        """
        Format an agent name with its assigned color.
        
        Args:
            agent_id: The agent's ID
            agent_name: The agent's display name
            
        Returns:
            Colored agent name string
        """
        if not self.enabled:
            return agent_name
        
        color = self._assign_color(agent_id)
        return f"{color}{agent_name}{Colors.RESET}"
    
    def print_agent_list(self, agents: list[tuple[str, str]]) -> None:
        """
        Print a list of agents with their assigned colors.
        
        Args:
            agents: List of (agent_id, agent_name) tuples
        """
        print("Agents:")
        for agent_id, agent_name in agents:
            formatted_name = self.format_agent_name(agent_id, agent_name)
            print(f"  • {formatted_name}")


def create_colored_printer(enabled: bool = True) -> ColoredAgentDisplay:
    """
    Create a colored display printer for workspace messages.
    
    This is a convenience function for creating a ColoredAgentDisplay
    instance that can be used as an on_message callback.
    
    Args:
        enabled: Whether to enable colors (disable for non-terminal output)
    
    Returns:
        ColoredAgentDisplay instance
    
    Example:
        ```python
        printer = synqed.create_colored_printer()
        
        workspace = synqed.create_workspace(
            name="My Workspace",
            task="Collaborate on a task",
            agents=[...],
            on_message=printer.print_message,
        )
        ```
    """
    return ColoredAgentDisplay(enabled=enabled)

