"""
Model router for OomLlama.

Routes queries to the most appropriate model based on:
- Query complexity
- Required capabilities
- Available resources
"""

from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any
from enum import Enum


class ModelCapability(Enum):
    """Model capabilities."""
    FAST = "fast"           # Quick responses
    REASONING = "reasoning"  # Complex logic
    CODE = "code"           # Code generation
    CREATIVE = "creative"   # Creative writing
    MULTILINGUAL = "multilingual"  # Multiple languages
    VISION = "vision"       # Image understanding


@dataclass
class ModelConfig:
    """Configuration for a model."""
    name: str
    size: str  # "3b", "7b", "32b", "70b"
    capabilities: List[ModelCapability] = field(default_factory=list)
    context_length: int = 4096
    priority: int = 0  # Higher = preferred
    ollama_url: str = "http://localhost:11434"

    @property
    def size_gb(self) -> float:
        """Approximate model size in GB."""
        sizes = {"3b": 2.0, "7b": 4.5, "14b": 9.0, "32b": 20.0, "70b": 40.0}
        return sizes.get(self.size.lower(), 5.0)


# Default model configurations
DEFAULT_MODELS = [
    ModelConfig(
        name="qwen2.5:3b",
        size="3b",
        capabilities=[ModelCapability.FAST, ModelCapability.MULTILINGUAL],
        context_length=32768,
        priority=10
    ),
    ModelConfig(
        name="qwen2.5:7b",
        size="7b",
        capabilities=[ModelCapability.FAST, ModelCapability.CODE, ModelCapability.MULTILINGUAL],
        context_length=32768,
        priority=20
    ),
    ModelConfig(
        name="qwen2.5:32b",
        size="32b",
        capabilities=[ModelCapability.REASONING, ModelCapability.CODE, ModelCapability.MULTILINGUAL],
        context_length=32768,
        priority=30
    ),
    ModelConfig(
        name="deepseek-coder:6.7b",
        size="7b",
        capabilities=[ModelCapability.CODE, ModelCapability.FAST],
        context_length=16384,
        priority=25
    ),
    ModelConfig(
        name="llama3.2:3b",
        size="3b",
        capabilities=[ModelCapability.FAST, ModelCapability.CREATIVE],
        context_length=8192,
        priority=15
    ),
]


class ModelRouter:
    """
    Routes queries to appropriate models.

    Example:
        router = ModelRouter()
        model = router.route("Write a Python function")
        print(model.name)  # "deepseek-coder:6.7b"
    """

    def __init__(self, models: Optional[List[ModelConfig]] = None):
        """
        Initialize router.

        Args:
            models: Available models (uses defaults if not specified)
        """
        self.models = models or DEFAULT_MODELS.copy()
        self._available_cache: Optional[List[str]] = None

    def add_model(self, model: ModelConfig):
        """Add a model configuration."""
        self.models.append(model)
        self._available_cache = None

    def route(
        self,
        query: str,
        prefer_fast: bool = False,
        require_capability: Optional[ModelCapability] = None,
        max_size: Optional[str] = None
    ) -> ModelConfig:
        """
        Route query to best model.

        Args:
            query: User query
            prefer_fast: Prioritize speed over quality
            require_capability: Required capability
            max_size: Maximum model size

        Returns:
            Best matching ModelConfig
        """
        candidates = self.models.copy()

        # Filter by capability
        if require_capability:
            candidates = [m for m in candidates if require_capability in m.capabilities]

        # Filter by size
        if max_size:
            size_order = ["3b", "7b", "14b", "32b", "70b"]
            max_idx = size_order.index(max_size.lower()) if max_size.lower() in size_order else len(size_order)
            candidates = [m for m in candidates if size_order.index(m.size.lower()) <= max_idx]

        if not candidates:
            # Fallback to default
            return self.models[0]

        # Auto-detect requirements
        query_lower = query.lower()

        # Code detection
        code_keywords = ["code", "function", "class", "python", "javascript", "def ", "import ", "```"]
        if any(kw in query_lower for kw in code_keywords):
            code_models = [m for m in candidates if ModelCapability.CODE in m.capabilities]
            if code_models:
                candidates = code_models

        # Complexity detection (longer = more complex)
        if len(query) > 500 and not prefer_fast:
            reasoning_models = [m for m in candidates if ModelCapability.REASONING in m.capabilities]
            if reasoning_models:
                candidates = reasoning_models

        # Speed preference
        if prefer_fast:
            fast_models = [m for m in candidates if ModelCapability.FAST in m.capabilities]
            if fast_models:
                candidates = fast_models

        # Sort by priority and return best
        candidates.sort(key=lambda m: m.priority, reverse=True)
        return candidates[0]

    def route_with_reason(
        self,
        query: str,
        **kwargs
    ) -> tuple[ModelConfig, str]:
        """
        Route with explanation.

        Returns:
            (model, reason) tuple
        """
        model = self.route(query, **kwargs)

        reasons = []
        if ModelCapability.CODE in model.capabilities and "code" in query.lower():
            reasons.append("code query detected")
        if ModelCapability.FAST in model.capabilities:
            reasons.append("fast response")
        if ModelCapability.REASONING in model.capabilities and len(query) > 500:
            reasons.append("complex query")

        reason = ", ".join(reasons) if reasons else f"default selection (priority={model.priority})"

        return model, reason

    def list_models(self) -> List[str]:
        """List configured model names."""
        return [m.name for m in self.models]
