"""
Address resolver — resolves agent names/roles to their network addresses.
"""

from __future__ import annotations
from typing import Dict, List, Optional


class AddressResolver:
    """Resolves symbolic agent names and roles to concrete agent IDs."""

    def __init__(self) -> None:
        self._name_table: Dict[str, str] = {}  # name -> agent_id
        self._role_table: Dict[str, List[str]] = {}  # role -> [agent_ids]

    def register(self, name: str, agent_id: str, role: Optional[str] = None) -> None:
        self._name_table[name] = agent_id
        if role:
            self._role_table.setdefault(role, []).append(agent_id)

    def resolve_by_name(self, name: str) -> Optional[str]:
        return self._name_table.get(name)

    def resolve_by_role(self, role: str) -> List[str]:
        return self._role_table.get(role, [])

    def unregister(self, name: str) -> Optional[str]:
        agent_id = self._name_table.pop(name, None)
        if agent_id:
            for role, ids in self._role_table.items():
                self._role_table[role] = [i for i in ids if i != agent_id]
        return agent_id