"""Parallel agent discovery — probe each MCP key and group by org."""

import asyncio
from typing import Dict, List, Tuple

from mercury_launcher.log import get_logger
from mercury_launcher.models import AgentConfig, AgentContext, DiscoveredAgent, OrgInfo
from mercury_launcher.poller import MercuryPoller

logger = get_logger("discover")


async def _probe_agent(agent_config: AgentConfig) -> Tuple[AgentConfig, AgentContext | None, str | None]:
    """Call mercury_get_agent_context for a single agent. Returns (config, context, error)."""
    poller = MercuryPoller(endpoint=agent_config.mercury_url, api_key=agent_config.mcp_api_key)
    try:
        raw = await poller.get_agent_context()
        context = AgentContext(**raw)
        return (agent_config, context, None)
    except Exception as e:
        return (agent_config, None, str(e))
    finally:
        await poller.close()


async def discover_agents(agents: List[AgentConfig]) -> Dict[str, OrgInfo]:
    """Probe all configured agents in parallel and group by org.

    Returns: {org_id: OrgInfo(org_name, agents=[...])}
    """
    results = await asyncio.gather(*[_probe_agent(a) for a in agents])

    orgs: Dict[str, OrgInfo] = {}
    for config, context, error in results:
        if error or context is None:
            logger.warning("  ✗ %s — %s", config.name, error or "no response")
            continue

        name = context.name or config.name
        role = context.agent_role or ""
        logger.info("  ✓ %s (%s)", name, role)

        org_id = context.org_id or "unknown"
        if org_id not in orgs:
            orgs[org_id] = OrgInfo(org_name=context.org_name or "Unknown")
        orgs[org_id].agents.append(DiscoveredAgent(config=config, context=context))

    return orgs
