"""
Task Router for Intent-Based Routing

Provides:
- Intent classification
- Agent selection
- Fallback handling
- Routing metrics
"""

from typing import Dict, Optional, List, Any, Callable, Awaitable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
import logging
import re

logger = logging.getLogger(__name__)


class RoutingStrategy(Enum):
    """Routing strategies"""
    DIRECT = "direct"           # Route to specific agent
    ROUND_ROBIN = "round_robin"  # Distribute evenly
    LEAST_LOADED = "least_loaded"  # Route to least loaded
    PRIORITY = "priority"        # Route based on priority


@dataclass
class RoutingRule:
    """A rule for routing tasks to agents"""
    intent_pattern: str
    agent_type: str
    priority: int = 0
    conditions: Dict[str, Any] = field(default_factory=dict)
    fallback_agent: Optional[str] = None

    def matches(self, intent: str, context: Dict[str, Any]) -> bool:
        """Check if this rule matches the intent and context"""
        # Check intent pattern (supports wildcards)
        pattern = self.intent_pattern.replace("*", ".*")
        if not re.match(f"^{pattern}$", intent, re.IGNORECASE):
            return False

        # Check conditions
        for key, value in self.conditions.items():
            if context.get(key) != value:
                return False

        return True


@dataclass
class RoutingResult:
    """Result of a routing decision"""
    success: bool
    agent_id: Optional[str] = None
    agent_type: Optional[str] = None
    is_fallback: bool = False
    reason: str = ""
    routing_time_ms: float = 0


class TaskRouter:
    """
    Routes tasks to appropriate agents based on intent.

    Features:
    - Pattern-based routing rules
    - Priority ordering
    - Fallback chains
    - Routing metrics
    """

    def __init__(self, supervisor=None):
        """
        Initialize the task router.

        Args:
            supervisor: AgentSupervisor for agent availability
        """
        self.supervisor = supervisor
        self.rules: List[RoutingRule] = []
        self._default_routes: Dict[str, str] = {}
        self._routing_history: List[Dict[str, Any]] = []
        self._max_history = 1000

        # Initialize default rules
        self._setup_default_rules()

    def _setup_default_rules(self) -> None:
        """Sets up default routing rules"""
        default_rules = [
            # Learning support
            RoutingRule(
                intent_pattern="learning_*",
                agent_type="tutor-agent",
                priority=10,
                fallback_agent="orchestrator"
            ),
            RoutingRule(
                intent_pattern="help_*",
                agent_type="tutor-agent",
                priority=5
            ),
            RoutingRule(
                intent_pattern="explain_*",
                agent_type="tutor-agent",
                priority=5
            ),

            # Grading
            RoutingRule(
                intent_pattern="grade_*",
                agent_type="grader-agent",
                priority=10,
                fallback_agent="quality-judge"
            ),
            RoutingRule(
                intent_pattern="evaluate_*",
                agent_type="grader-agent",
                priority=5
            ),
            RoutingRule(
                intent_pattern="correct_*",
                agent_type="grader-agent",
                priority=5
            ),

            # Quality checks
            RoutingRule(
                intent_pattern="quality_*",
                agent_type="quality-judge",
                priority=10
            ),
            RoutingRule(
                intent_pattern="review_*",
                agent_type="quality-judge",
                priority=5
            ),

            # Alerts
            RoutingRule(
                intent_pattern="alert_*",
                agent_type="alert-agent",
                priority=10
            ),
            RoutingRule(
                intent_pattern="notify_*",
                agent_type="alert-agent",
                priority=5
            ),

            # System/Admin
            RoutingRule(
                intent_pattern="system_*",
                agent_type="orchestrator",
                priority=10
            ),
            RoutingRule(
                intent_pattern="admin_*",
                agent_type="orchestrator",
                priority=10
            ),
        ]

        for rule in default_rules:
            self.add_rule(rule)

    def add_rule(self, rule: RoutingRule) -> None:
        """
        Adds a routing rule.

        Args:
            rule: The routing rule to add
        """
        self.rules.append(rule)
        # Sort by priority (higher first)
        self.rules.sort(key=lambda r: r.priority, reverse=True)

    def remove_rule(self, intent_pattern: str) -> bool:
        """
        Removes a routing rule by pattern.

        Args:
            intent_pattern: The pattern to remove

        Returns:
            True if a rule was removed
        """
        original_len = len(self.rules)
        self.rules = [r for r in self.rules if r.intent_pattern != intent_pattern]
        return len(self.rules) < original_len

    def set_default_route(self, agent_type: str, agent_id: str) -> None:
        """
        Sets a default agent for a type.

        Args:
            agent_type: The agent type
            agent_id: The default agent ID
        """
        self._default_routes[agent_type] = agent_id

    async def route(
        self,
        intent: str,
        context: Optional[Dict[str, Any]] = None,
        strategy: RoutingStrategy = RoutingStrategy.LEAST_LOADED
    ) -> RoutingResult:
        """
        Routes a task based on intent.

        Args:
            intent: The task intent
            context: Additional context for routing
            strategy: Load balancing strategy

        Returns:
            RoutingResult with agent assignment
        """
        start_time = datetime.now(timezone.utc)
        context = context or {}

        # Find matching rule
        matching_rule = None
        for rule in self.rules:
            if rule.matches(intent, context):
                matching_rule = rule
                break

        if not matching_rule:
            result = RoutingResult(
                success=False,
                reason=f"No routing rule matches intent: {intent}"
            )
            self._record_routing(intent, result)
            return result

        # Get available agent
        agent_id = await self._get_agent(
            matching_rule.agent_type,
            strategy
        )

        if agent_id:
            result = RoutingResult(
                success=True,
                agent_id=agent_id,
                agent_type=matching_rule.agent_type,
                is_fallback=False,
                reason="Primary agent selected"
            )
        elif matching_rule.fallback_agent:
            # Try fallback
            agent_id = await self._get_agent(
                matching_rule.fallback_agent,
                strategy
            )
            if agent_id:
                result = RoutingResult(
                    success=True,
                    agent_id=agent_id,
                    agent_type=matching_rule.fallback_agent,
                    is_fallback=True,
                    reason="Fallback agent selected"
                )
            else:
                result = RoutingResult(
                    success=False,
                    reason="No agents available (primary or fallback)"
                )
        else:
            result = RoutingResult(
                success=False,
                reason=f"No {matching_rule.agent_type} agents available"
            )

        # Calculate routing time
        end_time = datetime.now(timezone.utc)
        result.routing_time_ms = (end_time - start_time).total_seconds() * 1000

        self._record_routing(intent, result)

        return result

    async def _get_agent(
        self,
        agent_type: str,
        strategy: RoutingStrategy
    ) -> Optional[str]:
        """Gets an available agent of the specified type"""
        # Check default route first
        if agent_type in self._default_routes:
            agent_id = self._default_routes[agent_type]
            if self.supervisor and self.supervisor.agents.get(agent_id):
                info = self.supervisor.agents[agent_id]
                if info.is_available():
                    return agent_id

        # Use supervisor for load balancing
        if self.supervisor:
            strategy_str = "least_loaded"
            if strategy == RoutingStrategy.ROUND_ROBIN:
                strategy_str = "round_robin"

            return self.supervisor.get_available_agent(
                agent_type,
                strategy_str
            )

        return None

    def _record_routing(
        self,
        intent: str,
        result: RoutingResult
    ) -> None:
        """Records routing decision for analytics"""
        record = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "intent": intent,
            "success": result.success,
            "agent_id": result.agent_id,
            "agent_type": result.agent_type,
            "is_fallback": result.is_fallback,
            "routing_time_ms": result.routing_time_ms,
            "reason": result.reason
        }

        self._routing_history.append(record)

        # Trim history
        if len(self._routing_history) > self._max_history:
            self._routing_history = self._routing_history[-self._max_history:]

        # Log
        if result.success:
            logger.debug(
                f"Routed '{intent}' to {result.agent_id} "
                f"({'fallback' if result.is_fallback else 'primary'})"
            )
        else:
            logger.warning(f"Failed to route '{intent}': {result.reason}")

    # Analytics

    def get_routing_stats(self) -> Dict[str, Any]:
        """Gets routing statistics"""
        if not self._routing_history:
            return {
                "total_routes": 0,
                "success_rate": 0,
                "fallback_rate": 0,
                "avg_routing_time_ms": 0
            }

        total = len(self._routing_history)
        successful = sum(1 for r in self._routing_history if r["success"])
        fallbacks = sum(1 for r in self._routing_history if r["is_fallback"])
        avg_time = sum(
            r["routing_time_ms"] for r in self._routing_history
        ) / total

        return {
            "total_routes": total,
            "successful_routes": successful,
            "success_rate": successful / total,
            "fallback_routes": fallbacks,
            "fallback_rate": fallbacks / max(successful, 1),
            "avg_routing_time_ms": avg_time
        }

    def get_intent_distribution(self) -> Dict[str, int]:
        """Gets distribution of routed intents"""
        distribution: Dict[str, int] = {}
        for record in self._routing_history:
            intent = record["intent"]
            # Extract base intent (before _)
            base = intent.split("_")[0] if "_" in intent else intent
            distribution[base] = distribution.get(base, 0) + 1
        return distribution

    def get_agent_distribution(self) -> Dict[str, int]:
        """Gets distribution of routes by agent type"""
        distribution: Dict[str, int] = {}
        for record in self._routing_history:
            agent_type = record.get("agent_type", "unknown")
            if agent_type:
                distribution[agent_type] = distribution.get(agent_type, 0) + 1
        return distribution

    def get_failure_reasons(self) -> Dict[str, int]:
        """Gets distribution of routing failure reasons"""
        reasons: Dict[str, int] = {}
        for record in self._routing_history:
            if not record["success"]:
                reason = record["reason"]
                reasons[reason] = reasons.get(reason, 0) + 1
        return reasons

    def clear_history(self) -> None:
        """Clears routing history"""
        self._routing_history.clear()

    # Rule inspection

    def get_rules(self) -> List[Dict[str, Any]]:
        """Gets all routing rules as dicts"""
        return [
            {
                "intent_pattern": r.intent_pattern,
                "agent_type": r.agent_type,
                "priority": r.priority,
                "conditions": r.conditions,
                "fallback_agent": r.fallback_agent
            }
            for r in self.rules
        ]

    def find_matching_rules(
        self,
        intent: str,
        context: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """Finds all rules that match an intent"""
        context = context or {}
        return [
            {
                "intent_pattern": r.intent_pattern,
                "agent_type": r.agent_type,
                "priority": r.priority
            }
            for r in self.rules
            if r.matches(intent, context)
        ]
