"""
Task Router - Routes tasks to optimal agents based on multiple factors.

The router considers:
1. Task type and complexity
2. Agent capabilities and availability
3. Budget constraints
4. Risk level and approval requirements
5. Current agent health

Usage:
    from agent_orchestrator.routing.router import TaskRouter, RouteDecision

    router = TaskRouter(db, budget_manager)
    decision = await router.route("Write unit tests for auth module")

    if decision.approved:
        response = await decision.adapter.execute(task, context)
"""

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Optional

from .task_types import (
    TaskType,
    AgentTool,
    RoutingRule,
    AgentCapability,
    ROUTING_TABLE,
    AGENT_PROFILES,
    classify_task,
    get_routing_rule,
)
from ..adapters.base import BaseAdapter, CLIAgentAdapter, RiskLevel, AgentStatus
from ..persistence.database import OrchestratorDB
from ..config import Config, get_config


logger = logging.getLogger(__name__)


@dataclass
class RouteDecision:
    """The result of routing a task to an agent."""

    task: str
    task_type: TaskType
    selected_agent: AgentTool
    adapter: Optional[BaseAdapter]
    risk_level: RiskLevel
    autonomy_mode: Optional[str]
    approved: bool
    requires_approval: bool
    rejection_reason: Optional[str] = None
    estimated_cost: float = 0.0
    context_requirement: str = "normal"
    routing_notes: str = ""
    alternatives: list[AgentTool] = field(default_factory=list)
    timestamp: datetime = field(default_factory=datetime.now)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for logging/storage."""
        return {
            "task": self.task[:100],
            "task_type": self.task_type.name,
            "selected_agent": self.selected_agent.value,
            "risk_level": self.risk_level.value,
            "autonomy_mode": self.autonomy_mode,
            "approved": self.approved,
            "requires_approval": self.requires_approval,
            "rejection_reason": self.rejection_reason,
            "estimated_cost": self.estimated_cost,
            "alternatives": [a.value for a in self.alternatives],
            "timestamp": self.timestamp.isoformat(),
        }


class BudgetManager:
    """Manages budget constraints for agent usage."""

    def __init__(self, config: Config):
        self.config = config
        self.daily_spent: dict[str, float] = {}  # agent_id -> spent today
        self.task_spent: dict[str, float] = {}  # task_id -> spent on task

    def can_afford(self, agent: AgentTool, estimated_cost: float) -> bool:
        """Check if we can afford to run this agent."""
        agent_key = agent.value
        daily_limit = self.config.budgets.daily_budget_usd
        task_limit = self.config.budgets.task_budget_usd

        current_daily = self.daily_spent.get(agent_key, 0.0)

        # Check against daily limit
        if current_daily + estimated_cost > daily_limit:
            return False

        return True

    def record_spend(self, agent: AgentTool, task_id: str, cost: float) -> None:
        """Record spending for an agent/task."""
        agent_key = agent.value
        self.daily_spent[agent_key] = self.daily_spent.get(agent_key, 0.0) + cost
        self.task_spent[task_id] = self.task_spent.get(task_id, 0.0) + cost

    def get_remaining_budget(self, agent: AgentTool) -> float:
        """Get remaining daily budget for an agent."""
        agent_key = agent.value
        spent = self.daily_spent.get(agent_key, 0.0)
        return max(0, self.config.budgets.daily_budget_usd - spent)

    def estimate_cost(self, agent: AgentTool, context_size: str) -> float:
        """Estimate cost based on agent and context size."""
        profile = AGENT_PROFILES.get(agent)
        if not profile:
            return 0.0

        # Base cost estimates per context size (rough estimates)
        base_costs = {
            "minimal": 0.01,
            "normal": 0.05,
            "large": 0.15,
            "massive": 0.50,
        }

        base = base_costs.get(context_size, 0.05)

        # Adjust by agent cost tier
        tier_multiplier = {
            "low": 0.5,
            "medium": 1.0,
            "high": 2.0,
        }

        return base * tier_multiplier.get(profile.cost_tier, 1.0)


class TaskRouter:
    """
    Routes tasks to optimal agents based on multiple factors.

    The router:
    1. Classifies the task type
    2. Looks up routing rules
    3. Checks agent availability and health
    4. Validates budget constraints
    5. Applies risk/approval policies
    6. Returns a RouteDecision
    """

    def __init__(
        self,
        db: OrchestratorDB,
        config: Optional[Config] = None,
        adapters: Optional[dict[str, BaseAdapter]] = None,
    ):
        """
        Initialize the task router.

        Args:
            db: Database for persistence
            config: Configuration (uses default if not provided)
            adapters: Map of agent_id -> adapter instances
        """
        self.db = db
        self.config = config or get_config()
        self.adapters = adapters or {}
        self.budget_manager = BudgetManager(self.config)

        # Map AgentTool to registered adapters
        self._tool_to_adapter: dict[AgentTool, BaseAdapter] = {}

    def register_adapter(self, tool: AgentTool, adapter: BaseAdapter) -> None:
        """Register an adapter for an agent tool."""
        self._tool_to_adapter[tool] = adapter
        self.adapters[adapter.agent_id] = adapter

    def get_adapter(self, tool: AgentTool) -> Optional[BaseAdapter]:
        """Get the registered adapter for a tool."""
        return self._tool_to_adapter.get(tool)

    def _get_cli_budget(self, agent: AgentTool) -> Optional[Any]:
        """Get budget config for a CLI agent."""
        agent_map = {
            AgentTool.CLAUDE_CODE: "claude_code",
            AgentTool.GEMINI_CLI: "gemini_cli",
            AgentTool.CODEX_CLI: "codex",
        }
        budget_key = agent_map.get(agent)
        if not budget_key:
            return None
        return self.config.budgets.get_budget(budget_key)

    def _cli_limit_status(self, agent: AgentTool, adapter: BaseAdapter) -> tuple[str, str]:
        """Check if CLI usage is near or over its limit."""
        if not self.config.routing.cli_rebalance_enabled:
            return ("ok", "")

        if not isinstance(adapter, CLIAgentAdapter):
            return ("ok", "")

        budget = self._get_cli_budget(agent)
        if not budget:
            return ("ok", "")

        usage = self.db.get_daily_usage(adapter.agent_id)
        tokens_used = usage.get("tokens_input", 0) + usage.get("tokens_output", 0)
        cost_used = usage.get("cost_usd", 0.0)

        token_pct = (
            tokens_used / budget.daily_token_limit
            if budget.daily_token_limit > 0
            else 0.0
        )
        cost_pct = (
            cost_used / budget.daily_cost_limit
            if budget.daily_cost_limit > 0
            else 0.0
        )
        limit_pct = max(token_pct, cost_pct)

        if limit_pct >= self.config.routing.cli_hard_limit_pct:
            reason = f"usage {limit_pct:.0%} of daily limit"
            return ("blocked", reason)
        if limit_pct >= self.config.routing.cli_soft_limit_pct:
            reason = f"usage {limit_pct:.0%} of daily limit"
            return ("near", reason)

        return ("ok", "")

    async def route(
        self,
        task: str,
        context: Optional[dict[str, Any]] = None,
        force_agent: Optional[AgentTool] = None,
        task_type_override: Optional[TaskType] = None,
    ) -> RouteDecision:
        """
        Route a task to the optimal agent.

        Args:
            task: The task description
            context: Additional context (affects cost estimation)
            force_agent: Force routing to a specific agent
            task_type_override: Override automatic task classification

        Returns:
            RouteDecision with selected agent and approval status
        """
        context = context or {}

        # Step 1: Classify task
        if task_type_override:
            task_type = task_type_override
        else:
            task_type = classify_task(task)

        logger.debug(f"Classified task as: {task_type.name}")

        # Step 2: Get routing rule
        rule = get_routing_rule(task_type)

        # Step 3: If forcing agent, validate and return
        if force_agent:
            return await self._route_to_forced_agent(task, task_type, rule, force_agent, context)

        # Step 4: Find best available agent
        selected_agent = None
        alternatives = []
        near_limit_candidates: list[tuple[AgentTool, BaseAdapter, str]] = []
        rejection_reasons: list[str] = []

        for agent in rule.preferred_agents:
            # Check if adapter is registered
            adapter = self.get_adapter(agent)
            if not adapter:
                alternatives.append(agent)
                continue

            # Check authentication for CLI agents
            if (
                self.config.cli_auth.check_enabled
                and isinstance(adapter, CLIAgentAdapter)
                and not adapter.is_authenticated()
            ):
                reason = adapter.auth_error() or "CLI not authenticated"
                logger.debug(f"Agent {agent.value} not authenticated: {reason}")
                alternatives.append(agent)
                rejection_reasons.append(f"{agent.value} auth: {reason}")
                continue

            # Check health
            if not adapter.is_healthy():
                logger.debug(f"Agent {agent.value} is not healthy, skipping")
                alternatives.append(agent)
                continue

            # Check CLI limit usage for rebalancing
            limit_status, limit_reason = self._cli_limit_status(agent, adapter)
            if limit_status == "blocked":
                logger.debug(f"Agent {agent.value} blocked by limit: {limit_reason}")
                alternatives.append(agent)
                rejection_reasons.append(f"{agent.value} limit: {limit_reason}")
                continue
            if limit_status == "near":
                near_limit_candidates.append((agent, adapter, limit_reason))
                alternatives.append(agent)
                continue

            # Check budget
            estimated_cost = self.budget_manager.estimate_cost(agent, rule.context_requirement)
            if not self.budget_manager.can_afford(agent, estimated_cost):
                logger.debug(f"Agent {agent.value} exceeds budget, skipping")
                alternatives.append(agent)
                continue

            # Found a suitable agent
            selected_agent = agent
            break

        routing_notes = rule.notes
        if not selected_agent and near_limit_candidates:
            selected_agent, adapter, limit_reason = near_limit_candidates[0]
            routing_notes = f"Using near-limit agent {selected_agent.value}: {limit_reason}"

        if not selected_agent:
            rejection_reason = "No suitable agent available (check health, registration, budget)"
            if rejection_reasons:
                rejection_reason = "; ".join(rejection_reasons[:3])
            return RouteDecision(
                task=task,
                task_type=task_type,
                selected_agent=rule.preferred_agents[0],  # Default to first preference
                adapter=None,
                risk_level=rule.risk_level,
                autonomy_mode=rule.autonomy_recommendation,
                approved=False,
                requires_approval=True,
                rejection_reason=rejection_reason,
                alternatives=alternatives,
                context_requirement=rule.context_requirement,
                routing_notes=routing_notes,
            )

        # Step 5: Get adapter and prepare decision
        adapter = self.get_adapter(selected_agent)
        estimated_cost = self.budget_manager.estimate_cost(selected_agent, rule.context_requirement)

        # Step 6: Check approval requirements
        requires_approval, auto_approve = self._check_approval_requirements(rule.risk_level)

        return RouteDecision(
            task=task,
            task_type=task_type,
            selected_agent=selected_agent,
            adapter=adapter,
            risk_level=rule.risk_level,
            autonomy_mode=rule.autonomy_recommendation,
            approved=auto_approve,
            requires_approval=requires_approval,
            estimated_cost=estimated_cost,
            alternatives=[a for a in alternatives if a != selected_agent],
            context_requirement=rule.context_requirement,
            routing_notes=routing_notes,
        )

    async def _route_to_forced_agent(
        self,
        task: str,
        task_type: TaskType,
        rule: RoutingRule,
        forced_agent: AgentTool,
        context: dict[str, Any],
    ) -> RouteDecision:
        """Handle routing when agent is forced."""
        adapter = self.get_adapter(forced_agent)

        if not adapter:
            return RouteDecision(
                task=task,
                task_type=task_type,
                selected_agent=forced_agent,
                adapter=None,
                risk_level=rule.risk_level,
                autonomy_mode=rule.autonomy_recommendation,
                approved=False,
                requires_approval=True,
                rejection_reason=f"Forced agent {forced_agent.value} is not registered",
                context_requirement=rule.context_requirement,
            )

        if (
            self.config.cli_auth.check_enabled
            and isinstance(adapter, CLIAgentAdapter)
            and not adapter.is_authenticated()
        ):
            reason = adapter.auth_error() or "CLI not authenticated"
            return RouteDecision(
                task=task,
                task_type=task_type,
                selected_agent=forced_agent,
                adapter=None,
                risk_level=rule.risk_level,
                autonomy_mode=rule.autonomy_recommendation,
                approved=False,
                requires_approval=True,
                rejection_reason=f"Forced agent not authenticated: {reason}",
                context_requirement=rule.context_requirement,
            )

        if not adapter.is_healthy():
            return RouteDecision(
                task=task,
                task_type=task_type,
                selected_agent=forced_agent,
                adapter=adapter,
                risk_level=rule.risk_level,
                autonomy_mode=rule.autonomy_recommendation,
                approved=False,
                requires_approval=True,
                rejection_reason=f"Forced agent {forced_agent.value} is not healthy",
                context_requirement=rule.context_requirement,
            )

        limit_status, limit_reason = self._cli_limit_status(forced_agent, adapter)
        if limit_status == "blocked":
            return RouteDecision(
                task=task,
                task_type=task_type,
                selected_agent=forced_agent,
                adapter=None,
                risk_level=rule.risk_level,
                autonomy_mode=rule.autonomy_recommendation,
                approved=False,
                requires_approval=True,
                rejection_reason=f"Forced agent over limit: {limit_reason}",
                context_requirement=rule.context_requirement,
            )

        estimated_cost = self.budget_manager.estimate_cost(forced_agent, rule.context_requirement)
        requires_approval, auto_approve = self._check_approval_requirements(rule.risk_level)

        routing_notes = f"Forced to {forced_agent.value}"
        if limit_status == "near":
            routing_notes = f"{routing_notes} (near limit: {limit_reason})"

        return RouteDecision(
            task=task,
            task_type=task_type,
            selected_agent=forced_agent,
            adapter=adapter,
            risk_level=rule.risk_level,
            autonomy_mode=rule.autonomy_recommendation,
            approved=auto_approve,
            requires_approval=requires_approval,
            estimated_cost=estimated_cost,
            context_requirement=rule.context_requirement,
            routing_notes=routing_notes,
        )

    def _check_approval_requirements(self, risk_level: RiskLevel) -> tuple[bool, bool]:
        """
        Check if approval is required and if auto-approval is allowed.

        Returns:
            Tuple of (requires_approval, auto_approved)
        """
        # Risk-based approval policy
        if risk_level == RiskLevel.LOW:
            return False, True  # No approval needed, auto-approve

        elif risk_level == RiskLevel.MEDIUM:
            # Check config for medium risk auto-approval
            if self.config.control_loop.auto_approve_low_risk:
                return True, True  # Requires logging but auto-approved
            return True, False  # Needs manual approval

        elif risk_level == RiskLevel.HIGH:
            return True, False  # Always needs manual approval

        elif risk_level == RiskLevel.CRITICAL:
            return True, False  # Always needs manual approval (could auto-reject)

        return True, False  # Default to requiring approval

    async def execute_routed_task(
        self,
        decision: RouteDecision,
        context: Optional[dict[str, Any]] = None,
    ) -> Optional[Any]:
        """
        Execute a task that has been routed and approved.

        Args:
            decision: The routing decision
            context: Additional context for execution

        Returns:
            AgentResponse from the adapter, or None if not approved
        """
        if not decision.approved:
            logger.warning(f"Task not approved: {decision.rejection_reason}")
            return None

        if not decision.adapter:
            logger.error("No adapter available for approved task")
            return None

        context = context or {}

        # Set autonomy mode if applicable (for Codex)
        if decision.autonomy_mode and hasattr(decision.adapter, "set_autonomy_mode"):
            decision.adapter.set_autonomy_mode(decision.autonomy_mode)

        # Execute
        try:
            response = await decision.adapter.execute(decision.task, context)

            # Record cost
            if response.cost:
                self.budget_manager.record_spend(
                    decision.selected_agent,
                    f"task-{decision.timestamp.timestamp()}",
                    response.cost,
                )

            return response

        except Exception as e:
            logger.error(f"Error executing task: {e}")
            raise

    def get_available_agents(self) -> list[AgentTool]:
        """Get list of available (registered and healthy) agents."""
        available = []
        for tool, adapter in self._tool_to_adapter.items():
            if adapter.is_healthy():
                available.append(tool)
        return available

    def get_agent_status(self) -> dict[str, dict[str, Any]]:
        """Get status of all registered agents."""
        status = {}
        for tool, adapter in self._tool_to_adapter.items():
            status[tool.value] = {
                "healthy": adapter.is_healthy(),
                "status": adapter._status.value if hasattr(adapter, "_status") else "unknown",
                "remaining_budget": self.budget_manager.get_remaining_budget(tool),
            }
        return status


def create_router(
    db: OrchestratorDB,
    config: Optional[Config] = None,
) -> TaskRouter:
    """
    Factory function to create a configured TaskRouter.

    Args:
        db: Database instance
        config: Configuration (uses default if not provided)

    Returns:
        Configured TaskRouter
    """
    return TaskRouter(db=db, config=config)


async def route_and_execute(
    router: TaskRouter,
    task: str,
    context: Optional[dict[str, Any]] = None,
) -> Optional[Any]:
    """
    Convenience function to route and execute a task in one call.

    Args:
        router: The task router
        task: Task description
        context: Execution context

    Returns:
        AgentResponse or None
    """
    decision = await router.route(task, context)

    if not decision.approved:
        logger.info(f"Task requires approval: {decision.task_type.name} -> {decision.selected_agent.value}")
        return None

    return await router.execute_routed_task(decision, context)
