"""Agent Registry for a2a-mesh.

The registry is the service-discovery backbone of the mesh. Agents register
their Agent Cards, and the registry provides capability-based lookup, health
monitoring, and version management.
"""

from __future__ import annotations

import asyncio
from datetime import datetime, timezone
from typing import Sequence

import httpx

from a2a_mesh._logging import get_logger
from a2a_mesh.exceptions import (
    AgentAlreadyRegisteredError,
    AgentNotFoundError,
    HealthCheckFailedError,
)
from a2a_mesh.models import AgentCard, AgentStatus, RegisteredAgent

logger = get_logger(__name__)


class AgentRegistry:
    """In-memory agent registry with health monitoring.

    The registry stores agent cards, tracks health status, and supports
    capability-based discovery. A background health-check loop periodically
    pings registered agents and marks unhealthy ones.

    Attributes:
        agents: Mapping of agent name to its registered state.
        health_interval: Seconds between health check sweeps.
    """

    def __init__(self, health_interval: float = 30.0) -> None:
        """Initialize the registry.

        Args:
            health_interval: Seconds between health check sweeps.
        """
        self.agents: dict[str, RegisteredAgent] = {}
        self.health_interval = health_interval
        self._health_task: asyncio.Task[None] | None = None
        self._http_client: httpx.AsyncClient | None = None

    async def start(self) -> None:
        """Start the background health-check loop."""
        self._http_client = httpx.AsyncClient(timeout=5.0)
        self._health_task = asyncio.create_task(self._health_loop())
        logger.info("registry.started", health_interval=self.health_interval)

    async def stop(self) -> None:
        """Stop the health-check loop and clean up resources."""
        if self._health_task is not None:
            self._health_task.cancel()
            try:
                await self._health_task
            except asyncio.CancelledError:
                pass
            self._health_task = None
        if self._http_client is not None:
            await self._http_client.aclose()
            self._http_client = None
        logger.info("registry.stopped")

    def register(self, card: AgentCard, *, force: bool = False) -> RegisteredAgent:
        """Register an agent with the mesh.

        Args:
            card: The agent's capability card.
            force: If True, overwrite an existing registration.

        Returns:
            The registered agent record.

        Raises:
            AgentAlreadyRegisteredError: If the agent is already registered
                and force is False.
        """
        if card.name in self.agents and not force:
            raise AgentAlreadyRegisteredError(card.name)

        agent = RegisteredAgent(card=card, status=AgentStatus.UNKNOWN)
        self.agents[card.name] = agent
        logger.info(
            "agent.registered",
            agent=card.name,
            capabilities=card.capabilities,
            version=card.version,
        )
        return agent

    def deregister(self, agent_name: str) -> None:
        """Remove an agent from the registry.

        Args:
            agent_name: Name of the agent to remove.

        Raises:
            AgentNotFoundError: If the agent is not registered.
        """
        if agent_name not in self.agents:
            raise AgentNotFoundError(agent_name)
        del self.agents[agent_name]
        logger.info("agent.deregistered", agent=agent_name)

    def get(self, agent_name: str) -> RegisteredAgent:
        """Retrieve a registered agent by name.

        Args:
            agent_name: Name of the agent.

        Returns:
            The registered agent record.

        Raises:
            AgentNotFoundError: If the agent is not registered.
        """
        if agent_name not in self.agents:
            raise AgentNotFoundError(agent_name)
        return self.agents[agent_name]

    def find_by_capability(
        self,
        capabilities: Sequence[str],
        *,
        healthy_only: bool = True,
    ) -> list[RegisteredAgent]:
        """Find agents that support all of the given capabilities.

        Args:
            capabilities: Required capability tags.
            healthy_only: If True, exclude unhealthy agents.

        Returns:
            List of matching registered agents, sorted by current load
            (ascending).
        """
        required = set(capabilities)
        matches: list[RegisteredAgent] = []

        for agent in self.agents.values():
            agent_caps = set(agent.card.capabilities)
            if not required.issubset(agent_caps):
                continue
            if healthy_only and agent.status == AgentStatus.UNHEALTHY:
                continue
            matches.append(agent)

        matches.sort(key=lambda a: a.current_load)
        return matches

    def list_agents(self) -> list[RegisteredAgent]:
        """Return all registered agents.

        Returns:
            List of all registered agents.
        """
        return list(self.agents.values())

    async def check_health(self, agent_name: str) -> AgentStatus:
        """Run a health check against a single agent.

        Args:
            agent_name: Name of the agent to check.

        Returns:
            The updated health status.

        Raises:
            AgentNotFoundError: If the agent is not registered.
        """
        agent = self.get(agent_name)
        card = agent.card
        if not card.url:
            agent.status = AgentStatus.UNKNOWN
            return agent.status

        url = card.url.rstrip("/") + card.health_endpoint
        try:
            client = self._http_client or httpx.AsyncClient(timeout=5.0)
            resp = await client.get(url)
            if resp.status_code == 200:
                agent.status = AgentStatus.HEALTHY
            elif resp.status_code < 500:
                agent.status = AgentStatus.DEGRADED
            else:
                agent.status = AgentStatus.UNHEALTHY
        except httpx.HTTPError:
            agent.status = AgentStatus.UNHEALTHY

        agent.last_health_check = datetime.now(timezone.utc)
        logger.debug(
            "health.checked",
            agent=agent_name,
            status=agent.status.value,
        )
        return agent.status

    async def _health_loop(self) -> None:
        """Periodically check health of all registered agents."""
        while True:
            await asyncio.sleep(self.health_interval)
            for name in list(self.agents.keys()):
                try:
                    await self.check_health(name)
                except (AgentNotFoundError, HealthCheckFailedError):
                    pass
                except Exception:
                    logger.exception("health.loop_error", agent=name)
