"""Smart task router for a2a-mesh.

Routes incoming tasks to the best available agent based on capability
matching, load balancing, cost optimization, and custom routing policies.
"""

from __future__ import annotations

import random
from typing import Sequence

from a2a_mesh._logging import get_logger
from a2a_mesh.exceptions import NoCapableAgentError, QueueFullError
from a2a_mesh.models import (
    RegisteredAgent,
    RoutingPolicy,
    RoutingStrategy,
    Task,
)
from a2a_mesh.registry import AgentRegistry

logger = get_logger(__name__)


class Router:
    """Routes tasks to agents based on configurable policies.

    The router evaluates registered agents against the task's requirements
    and the active routing policy to select the optimal target.

    Attributes:
        registry: The agent registry to query.
        policy: The active routing policy.
    """

    def __init__(
        self,
        registry: AgentRegistry,
        policy: RoutingPolicy | None = None,
    ) -> None:
        """Initialize the router.

        Args:
            registry: The agent registry for discovery.
            policy: Routing policy configuration. Defaults to round-robin.
        """
        self.registry = registry
        self.policy = policy or RoutingPolicy()
        self._round_robin_index: int = 0

    def route(self, task: Task) -> RegisteredAgent:
        """Select the best agent for a task.

        Evaluates all registered agents against the task's requirements
        and the routing policy, then returns the optimal match.

        Args:
            task: The task to route.

        Returns:
            The selected agent.

        Raises:
            NoCapableAgentError: If no agent can handle the task.
            QueueFullError: If all capable agents are at capacity.
        """
        candidates = self._find_candidates(task)

        if not candidates:
            caps = task.required_capabilities or [task.agent]
            raise NoCapableAgentError(caps)

        # Filter out agents that are at capacity
        available = [
            a for a in candidates
            if a.current_load < self.policy.max_queue_depth
        ]

        if not available:
            first = candidates[0]
            raise QueueFullError(first.card.name, first.current_load)

        selected = self._apply_strategy(available)

        logger.info(
            "task.routed",
            task_id=task.task_id,
            agent=selected.card.name,
            strategy=self.policy.strategy.value,
            load=selected.current_load,
        )
        return selected

    def route_multi(
        self,
        task: Task,
        count: int,
    ) -> list[RegisteredAgent]:
        """Select multiple agents for fan-out execution.

        Args:
            task: The task to route.
            count: Number of agents to select.

        Returns:
            List of selected agents (up to count).

        Raises:
            NoCapableAgentError: If no agent can handle the task.
        """
        candidates = self._find_candidates(task)
        if not candidates:
            caps = task.required_capabilities or [task.agent]
            raise NoCapableAgentError(caps)

        available = [
            a for a in candidates
            if a.current_load < self.policy.max_queue_depth
        ]

        if not available:
            available = candidates

        # Return up to `count` agents, prioritised by strategy
        selected = self._sort_by_strategy(available)[:count]
        logger.info(
            "task.routed_multi",
            task_id=task.task_id,
            count=len(selected),
            agents=[a.card.name for a in selected],
        )
        return selected

    def _find_candidates(self, task: Task) -> list[RegisteredAgent]:
        """Find agents capable of handling the task.

        If the task specifies an agent name, look it up directly.
        Otherwise, search by required capabilities.
        """
        if task.agent:
            try:
                agent = self.registry.get(task.agent)
                return [agent]
            except Exception:
                return []

        if task.required_capabilities:
            return self.registry.find_by_capability(
                task.required_capabilities,
                healthy_only=True,
            )

        # No constraints: return all healthy agents
        return self.registry.find_by_capability([], healthy_only=True)

    def _apply_strategy(
        self,
        agents: Sequence[RegisteredAgent],
    ) -> RegisteredAgent:
        """Apply the routing strategy to select one agent."""
        strategy = self.policy.strategy

        if strategy == RoutingStrategy.ROUND_ROBIN:
            idx = self._round_robin_index % len(agents)
            self._round_robin_index += 1
            return agents[idx]

        if strategy == RoutingStrategy.LEAST_LOAD:
            return min(agents, key=lambda a: a.current_load)

        if strategy == RoutingStrategy.LEAST_COST:
            return min(agents, key=lambda a: a.card.cost_per_task)

        if strategy == RoutingStrategy.LEAST_LATENCY:
            return min(agents, key=lambda a: a.avg_latency_ms)

        if strategy == RoutingStrategy.RANDOM:
            return random.choice(list(agents))

        # Default fallback
        return agents[0]

    def _sort_by_strategy(
        self,
        agents: Sequence[RegisteredAgent],
    ) -> list[RegisteredAgent]:
        """Sort agents according to the active strategy."""
        strategy = self.policy.strategy

        if strategy == RoutingStrategy.LEAST_LOAD:
            return sorted(agents, key=lambda a: a.current_load)

        if strategy == RoutingStrategy.LEAST_COST:
            return sorted(agents, key=lambda a: a.card.cost_per_task)

        if strategy == RoutingStrategy.LEAST_LATENCY:
            return sorted(agents, key=lambda a: a.avg_latency_ms)

        return list(agents)
