"""
Flow Display for PraisonAI Agents

Visual display with agents in center and tools on sides.
"""

from typing import Dict, List, Set, Tuple
from collections import defaultdict
import threading
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.align import Align
from rich.columns import Columns
from rich.table import Table
from rich import box

class FlowDisplay:
    """Displays agent workflow with agents centered and tools on sides."""
    
    def __init__(self):
        self.console = Console()
        self.agents = []  # List of agents in order
        self.agent_tools = defaultdict(list)  # agent -> [tools]
        self.tracking = False
        self.lock = threading.Lock()
        
    def start(self):
        """Start tracking workflow."""
        self.tracking = True
        self.agents.clear()
        self.agent_tools.clear()
        
        # Register callbacks
        try:
            from praisonaiagents.main import register_display_callback
            
            def on_interaction(**kwargs):
                if self.tracking:
                    agent = kwargs.get('agent_name', 'Unknown')
                    self._add_agent(agent)
                    
            def on_tool_call(message, **kwargs):
                if self.tracking and "called function" in message:
                    parts = message.split("'")
                    if len(parts) > 1:
                        tool_name = parts[1]
                        agent = kwargs.get('agent_name', 'Unknown')
                        self._add_tool(agent, tool_name)
                        
            register_display_callback('interaction', on_interaction)
            register_display_callback('tool_call', on_tool_call)
            
        except ImportError:
            pass
            
    def stop(self):
        """Stop tracking and display the flow."""
        self.tracking = False
        self.display()
        
    def _add_agent(self, name: str):
        """Add an agent if not already present."""
        with self.lock:
            if name not in self.agents:
                self.agents.append(name)
                
    def _add_tool(self, agent_name: str, tool_name: str):
        """Add a tool to an agent."""
        with self.lock:
            if agent_name not in self.agents:
                self.agents.append(agent_name)
            if tool_name not in self.agent_tools[agent_name]:
                self.agent_tools[agent_name].append(tool_name)
                
    def display(self):
        """Display the flow chart with agents in center and tools on sides."""
        if not self.agents:
            return
            
        self.console.print("\n[bold cyan]🔄 Agent Workflow Flow[/bold cyan]\n")
        
        # Display start
        self._display_centered_node("── start ──", "grey35")
        self._display_arrow_down()
        
        # Display each agent with their tools
        for i, agent in enumerate(self.agents):
            self._display_agent_with_tools(agent)
            
            # Add arrow to next agent or end
            if i < len(self.agents) - 1:
                self._display_arrow_down()
        
        # Display end
        self._display_arrow_down()
        self._display_centered_node("── end ──", "grey35")
        
    def _display_agent_with_tools(self, agent: str):
        """Display agent with tools on the sides."""
        tools = self.agent_tools.get(agent, [])
        
        if not tools:
            # No tools - just agent
            self._display_centered_node(agent, "purple")
        else:
            # Split tools between left and right
            left_tools = tools[::2]  # Even indices
            right_tools = tools[1::2]  # Odd indices
            
            # Create the layout
            table = Table(show_header=False, show_edge=False, box=None, padding=0)
            table.add_column(justify="center", min_width=20)  # Left tools
            table.add_column(justify="center", min_width=5)   # Space
            table.add_column(justify="center", min_width=20)  # Agent
            table.add_column(justify="center", min_width=5)   # Space
            table.add_column(justify="center", min_width=20)  # Right tools
            
            # Create panels
            left_panel = self._create_tools_panel(left_tools) if left_tools else ""
            agent_panel = Panel(
                Text(agent, style="white on purple", justify="center"),
                style="white on purple",
                box=box.ROUNDED,
                padding=(0, 2)
            )
            right_panel = self._create_tools_panel(right_tools) if right_tools else ""
            
            # Add row
            table.add_row(left_panel, "", agent_panel, "", right_panel)
            
            # Display centered
            self.console.print(Align.center(table))
            
            # Show arrows
            if left_tools or right_tools:
                arrow_parts = []
                if left_tools:
                    arrow_parts.append("←→")
                else:
                    arrow_parts.append("  ")
                    
                arrow_parts.append("        ")  # Center space
                
                if right_tools:
                    arrow_parts.append("←→")
                else:
                    arrow_parts.append("  ")
                    
                self.console.print(Align.center(Text("".join(arrow_parts))))
                
    def _create_tools_panel(self, tools: List[str]) -> Panel:
        """Create a panel for tools."""
        if not tools:
            return ""
            
        if len(tools) == 1:
            return Panel(
                Text(tools[0], style="black on yellow", justify="center"),
                style="black on yellow",
                box=box.ROUNDED,
                padding=(0, 1)
            )
        else:
            # Multiple tools
            content = "\n".join(tools)
            return Panel(
                Text(content, style="black on yellow", justify="center"),
                style="black on yellow",
                box=box.ROUNDED,
                padding=(0, 1)
            )
            
    def _display_centered_node(self, label: str, color: str):
        """Display a centered node."""
        panel = Panel(
            Text(label, style=f"white on {color}", justify="center"),
            style=f"white on {color}",
            box=box.ROUNDED,
            padding=(0, 2)
        )
        self.console.print(Align.center(panel))
        
    def _display_arrow_down(self):
        """Display a downward arrow."""
        self.console.print()
        self.console.print(Align.center("↓"))
        self.console.print()


# Simple function to create and use
def track_workflow():
    """Create a flow display tracker."""
    return FlowDisplay()