"""
MessageRouter - Routes messages between Agents.

This module provides the MessageRouter class that handles turn-based
routing of messages between agents using the inbox-based pattern.
"""

import httpx
from typing import Dict, List, Optional
from datetime import datetime

from synqed.memory import InboxMessage
from synqed.server import AgentResponse


class MessageRouter:
    """
    Routes messages between Agents using turn-based communication.
    
    This router manages the flow of messages between agents:
    - Sends messages to agent inboxes
    - Retrieves responses from agents
    - Routes responses to target agents
    - Manages turn-based conversation cycles
    
    Example:
        ```python
        router = MessageRouter()
        
        router.register_agent("Writer", "http://localhost:8001")
        router.register_agent("Editor", "http://localhost:8002")
        
        # Start a conversation
        await router.send_message("USER", "Writer", "Write a story")
        
        # Run turn-based conversation
        transcript = await router.start_turn_cycle(
            initial_agent="Writer",
            num_turns=6
        )
        ```
    """
    
    def __init__(self, timeout: float = 30.0):
        """
        Initialize the message router.
        
        Args:
            timeout: HTTP request timeout in seconds (default: 30.0)
        """
        self.timeout = timeout
        self._agent_urls: Dict[str, str] = {}  # agent_name -> url
        self._transcript: List[Dict[str, str]] = []
    
    def register_agent(self, agent_name: str, agent_url: str) -> None:
        """
        Register an agent with the router.
        
        Args:
            agent_name: Name of the agent (must match send_to values)
            agent_url: Base URL of the agent server
        """
        self._agent_urls[agent_name] = agent_url.rstrip("/")
    
    def unregister_agent(self, agent_name: str) -> None:
        """
        Unregister an agent.
        
        Args:
            agent_name: Name of the agent to unregister
        """
        if agent_name in self._agent_urls:
            del self._agent_urls[agent_name]
    
    def list_agents(self) -> List[str]:
        """
        List all registered agent names.
        
        Returns:
            List of registered agent names
        """
        return list(self._agent_urls.keys())
    
    async def send_message(
        self,
        from_agent: str,
        to_agent: str,
        content: str
    ) -> dict:
        """
        Send a message to an agent's inbox.
        
        Args:
            from_agent: Name of the sending agent (or "USER" for user messages)
            to_agent: Name of the target agent
            content: Message content
            
        Returns:
            Response dictionary from the inbox endpoint
            
        Raises:
            ValueError: If target agent is not registered
            httpx.HTTPError: If HTTP request fails
        """
        if to_agent not in self._agent_urls:
            raise ValueError(f"Agent '{to_agent}' is not registered")
        
        agent_url = self._agent_urls[to_agent]
        
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            response = await client.post(
                f"{agent_url}/inbox",
                json={
                    "from_agent": from_agent,
                    "content": content
                }
            )
            response.raise_for_status()
            return response.json()
    
    async def get_agent_response(self, agent_name: str) -> AgentResponse:
        """
        Get a response from an agent by calling /respond.
        
        Args:
            agent_name: Name of the agent to get response from
            
        Returns:
            AgentResponse with send_to and content
            
        Raises:
            ValueError: If agent is not registered
            httpx.HTTPError: If HTTP request fails
        """
        if agent_name not in self._agent_urls:
            raise ValueError(f"Agent '{agent_name}' is not registered")
        
        agent_url = self._agent_urls[agent_name]
        
        async with httpx.AsyncClient(timeout=self.timeout) as client:
            response = await client.post(f"{agent_url}/respond")
            response.raise_for_status()
            return AgentResponse(**response.json())
    
    async def route_message(
        self,
        from_agent: str,
        send_to: str,
        content: str
    ) -> dict:
        """
        Route a message from one agent to another.
        
        This is a convenience method that combines send_message
        with automatic transcript recording.
        
        Args:
            from_agent: Name of the sending agent
            send_to: Name of the target agent
            content: Message content
            
        Returns:
            Response dictionary from the inbox endpoint
        """
        result = await self.send_message(from_agent, send_to, content)
        
        # Record in transcript
        self._transcript.append({
            "from": from_agent,
            "to": send_to,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })
        
        return result
    
    async def start_turn_cycle(
        self,
        initial_agent: str,
        initial_message: Optional[str] = None,
        num_turns: int = 6,
        initial_from: str = "USER"
    ) -> List[Dict[str, str]]:
        """
        Run a turn-based conversation cycle.
        
        This method:
        1. Sends initial message to the starting agent (if provided)
        2. Gets response from current agent
        3. Routes response to target agent
        4. Repeats for specified number of turns
        
        Args:
            initial_agent: Name of the agent to start with
            initial_message: Optional initial message to send
            num_turns: Number of conversation turns to execute
            initial_from: Name of the sender for initial message (default: "USER")
            
        Returns:
            List of message dictionaries from the transcript
            
        Raises:
            ValueError: If initial agent is not registered
        """
        if initial_agent not in self._agent_urls:
            raise ValueError(f"Initial agent '{initial_agent}' is not registered")
        
        # Clear transcript for new cycle
        self._transcript = []
        
        # Send initial message if provided
        if initial_message:
            await self.route_message(initial_from, initial_agent, initial_message)
        
        # Track current speaker
        current_agent = initial_agent
        
        # Run turns
        for turn in range(num_turns):
            try:
                # Get response from current agent
                response = await self.get_agent_response(current_agent)
                
                # Route message to target agent's inbox
                await self.route_message(
                    from_agent=current_agent,
                    send_to=response.send_to,
                    content=response.content
                )
                
                # Switch to target agent for next turn
                current_agent = response.send_to
                
            except Exception as e:
                # Log error and break
                self._transcript.append({
                    "from": "SYSTEM",
                    "to": "ERROR",
                    "content": f"Error on turn {turn + 1}: {str(e)}",
                    "timestamp": datetime.now().isoformat()
                })
                break
        
        return self._transcript
    
    def get_transcript(self) -> List[Dict[str, str]]:
        """
        Get the conversation transcript.
        
        Returns:
            List of message dictionaries
        """
        return self._transcript.copy()
    
    def clear_transcript(self) -> None:
        """Clear the conversation transcript."""
        self._transcript = []
    
    def __repr__(self) -> str:
        """String representation."""
        return (
            f"MessageRouter(agents={len(self._agent_urls)}, "
            f"transcript_length={len(self._transcript)})"
        )

