"""Smart model router — route LLM calls to the best model for each task type."""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import StrEnum

logger = logging.getLogger(__name__)


class TaskComplexity(StrEnum):
    """Classification of task complexity for routing decisions."""

    SIMPLE = "simple"
    MODERATE = "moderate"
    COMPLEX = "complex"


class TaskCategory(StrEnum):
    """Task categories that influence model selection."""

    LOOKUP = "lookup"
    SUMMARIZATION = "summarization"
    REASONING = "reasoning"
    CODING = "coding"
    CONVERSATION = "conversation"
    TOOL_USE = "tool_use"


@dataclass
class ModelTier:
    """A named model with its tier metadata."""

    model: str
    tier: str  # "fast", "balanced", "best"
    description: str = ""


@dataclass
class RoutingRule:
    """A single routing rule: match criteria → model selection."""

    category: TaskCategory | None = None
    complexity: TaskComplexity | None = None
    model: str = ""  # model to route to

    def matches(self, category: TaskCategory, complexity: TaskComplexity) -> bool:
        """Check if this rule matches the given category and complexity."""
        if self.category is not None and self.category != category:
            return False
        return not (self.complexity is not None and self.complexity != complexity)


@dataclass
class RouterConfig:
    """Configuration for the smart model router."""

    # Model tiers — users can customise these
    fast_model: str = "anthropic/claude-haiku-3"
    balanced_model: str = "anthropic/claude-sonnet-4-20250514"
    best_model: str = "anthropic/claude-sonnet-4-20250514"

    # Custom routing rules (evaluated first-match)
    rules: list[RoutingRule] = field(default_factory=list)

    # Whether routing is enabled (False → always use default model)
    enabled: bool = True


# Keywords used for heuristic task classification
_SIMPLE_KEYWORDS = frozenset(
    {
        "lookup",
        "find",
        "search",
        "what is",
        "who is",
        "define",
        "list",
        "show",
        "get",
        "fetch",
        "check",
    }
)
_COMPLEX_KEYWORDS = frozenset(
    {
        "analyse",
        "analyze",
        "compare",
        "plan",
        "design",
        "evaluate",
        "recommend",
        "explain why",
        "trade-off",
        "tradeoff",
        "pros and cons",
        "debug",
        "investigate",
        "optimise",
        "optimize",
        "refactor",
    }
)
_SUMMARY_KEYWORDS = frozenset(
    {
        "summarise",
        "summarize",
        "summary",
        "tldr",
        "brief",
        "recap",
        "overview",
        "digest",
    }
)
_CODING_KEYWORDS = frozenset(
    {
        "code",
        "implement",
        "function",
        "class",
        "script",
        "program",
        "algorithm",
        "fix bug",
        "write",
    }
)


def classify_task(
    task: str,
    step_count: int = 0,
    tool_count: int = 0,
) -> tuple[TaskCategory, TaskComplexity]:
    """Classify a task by category and complexity using heuristics.

    Uses keyword matching on the task description plus contextual signals
    like number of steps already taken and number of available tools.
    """
    lower = task.lower()

    # --- Category ---
    category = TaskCategory.CONVERSATION  # default

    if any(kw in lower for kw in _SUMMARY_KEYWORDS):
        category = TaskCategory.SUMMARIZATION
    elif any(kw in lower for kw in _CODING_KEYWORDS):
        category = TaskCategory.CODING
    elif any(kw in lower for kw in _SIMPLE_KEYWORDS):
        category = TaskCategory.LOOKUP
    elif any(kw in lower for kw in _COMPLEX_KEYWORDS):
        category = TaskCategory.REASONING

    if tool_count > 0 and category == TaskCategory.CONVERSATION:
        category = TaskCategory.TOOL_USE

    # --- Complexity ---
    complexity = TaskComplexity.MODERATE  # default

    if category == TaskCategory.LOOKUP:
        complexity = TaskComplexity.SIMPLE
    elif category in (TaskCategory.REASONING, TaskCategory.CODING):
        complexity = TaskComplexity.COMPLEX
    elif category == TaskCategory.SUMMARIZATION:
        complexity = TaskComplexity.MODERATE

    # Contextual signals can bump complexity
    if step_count >= 5:
        complexity = TaskComplexity.COMPLEX
    elif step_count == 0 and complexity == TaskComplexity.MODERATE:
        # First step with moderate task → keep moderate
        pass

    return category, complexity


# Default routing table: (complexity) → tier
_DEFAULT_ROUTING: dict[TaskComplexity, str] = {
    TaskComplexity.SIMPLE: "fast",
    TaskComplexity.MODERATE: "balanced",
    TaskComplexity.COMPLEX: "best",
}


class ModelRouter:
    """Routes LLM calls to the most appropriate model.

    Uses task classification + configurable rules to pick the right
    model tier (fast/balanced/best) for each call.
    """

    def __init__(self, config: RouterConfig | None = None) -> None:
        self._config = config or RouterConfig()

    @property
    def config(self) -> RouterConfig:
        return self._config

    def route(
        self,
        task: str,
        step_count: int = 0,
        tool_count: int = 0,
        override_model: str | None = None,
    ) -> str:
        """Select the best model for the given task context.

        Args:
            task: The task description.
            step_count: How many ReAct steps have been taken.
            tool_count: Number of available tools.
            override_model: Explicit model override (skips routing).

        Returns:
            The model string to use with LiteLLM.
        """
        # Explicit override always wins
        if override_model:
            return override_model

        # Routing disabled → balanced model
        if not self._config.enabled:
            return self._config.balanced_model

        category, complexity = classify_task(task, step_count, tool_count)

        # Check custom rules first (first match wins)
        for rule in self._config.rules:
            if rule.matches(category, complexity) and rule.model:
                logger.debug(
                    "Custom rule matched: %s/%s → %s",
                    category,
                    complexity,
                    rule.model,
                )
                return rule.model

        # Default routing by complexity → tier
        tier = _DEFAULT_ROUTING.get(complexity, "balanced")
        model = self._tier_to_model(tier)

        logger.debug(
            "Routed task (category=%s, complexity=%s) → %s (%s)",
            category,
            complexity,
            model,
            tier,
        )
        return model

    def _tier_to_model(self, tier: str) -> str:
        if tier == "fast":
            return self._config.fast_model
        if tier == "best":
            return self._config.best_model
        return self._config.balanced_model
