"""
Message router — directs messages to the correct recipient based on routing rules.
"""

from __future__ import annotations
from typing import Any, Dict, List, Optional
from .protocol import Message


class MessageRouter:
    """Routes messages between agents based on content, role, or address."""

    def __init__(self) -> None:
        self._routes: Dict[str, str] = {}  # routing_key -> agent_id
        self._fallback: Optional[str] = None

    def register(self, routing_key: str, agent_id: str) -> None:
        self._routes[routing_key] = agent_id

    def unregister(self, routing_key: str) -> Optional[str]:
        return self._routes.pop(routing_key, None)

    def route(self, message: Message) -> Optional[str]:
        payload_str = str(message.payload)
        for key, agent_id in self._routes.items():
            if key in payload_str:
                return agent_id
        return self._fallback

    def set_fallback(self, agent_id: str) -> None:
        self._fallback = agent_id

    def get_registered_routes(self) -> Dict[str, str]:
        return dict(self._routes)