"""Model router — selects the optimal model using learned embeddings."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path

from backend.config import EcateConfig, ModelConfig, ProviderConfig

logger = logging.getLogger("ecate.proxy.router")


@dataclass
class RouteDecision:
    """Result of routing decision."""
    provider: ProviderConfig
    model: ModelConfig
    score: float  # Embedding similarity score
    adjusted_score: float  # Score after strategy adjustment (cost-weighted)
    cost_per_mtok: float
    strategy: str


class Router:
    """Selects the best model for a request using learned embedding vectors.
    
    Each model has a learned 512d vector. Routing works by:
    1. Embed the prompt locally (~10ms)
    2. Dot product with each model's vector → scores
    3. Apply strategy to select model:
       - auto: score - λ*log(cost) — best tradeoff
       - best: highest score regardless of cost
       - cheap: cheapest model above threshold
    
    Falls back to simple routing (first matching model) if vectors not available.
    """

    def __init__(self, config: EcateConfig):
        self._config = config
        self._vector_router = None
        self._initialized = False
    
    def _ensure_initialized(self) -> None:
        """Lazy-initialize the vector router."""
        if self._initialized:
            return
        
        self._initialized = True
        
        try:
            from backend.router.vector_router import VectorRouter
            self._vector_router = VectorRouter(
                self._config,
                vectors_path=Path(self._config.routing.vectors_path),
                cost_weight=self._config.routing.cost_weight,
                score_threshold=self._config.routing.score_threshold,
            )
            if self._vector_router.has_vectors():
                models = self._vector_router.get_calibrated_models()
                logger.info(f"Router initialized with {len(models)} calibrated models")
            else:
                logger.info("Router initialized, no vectors loaded (using fallback)")
        except ImportError as e:
            logger.warning(f"Vector router not available: {e}")
            logger.warning("Using fallback routing. Install: pip install ecate-llm[router]")
    
    def has_vectors(self) -> bool:
        """Check if model vectors are loaded."""
        self._ensure_initialized()
        return self._vector_router is not None and self._vector_router.has_vectors()
    
    def get_calibrated_models(self) -> list[str]:
        """Get list of model IDs that have learned vectors."""
        self._ensure_initialized()
        if self._vector_router is None:
            return []
        return self._vector_router.get_calibrated_models()

    def route(
        self,
        prompt: str,
        strategy: str = "auto",
        needs_tools: bool = False,
        needs_vision: bool = False,
        token_estimate: int = 0,
        allowed_models: list[str] | None = None,
        exclude_models: set[str] | None = None,
        language: str = "en",
    ) -> RouteDecision | None:
        """Route prompt to optimal model.
        
        Args:
            prompt: User prompt to route
            strategy: Routing strategy
                - "auto": Balance quality vs cost (score - λ*log(cost))
                - "best": Highest score regardless of cost
                - "cheap": Cheapest model above score threshold
            needs_tools: Filter to models supporting tools
            needs_vision: Filter to models supporting vision
            token_estimate: Estimated total tokens (for context limit filter)
            allowed_models: Whitelist of model IDs (None = all)
            exclude_models: Blacklist of model IDs
            language: Detected language code (e.g., "en", "fr") or None
        
        Returns:
            RouteDecision or None if no suitable model found
        """
        self._ensure_initialized()
        
        # Try vector routing if available
        if self._vector_router is not None and self._vector_router.has_vectors():
            result = self._vector_router.route(
                prompt=prompt,
                strategy=strategy,
                needs_tools=needs_tools,
                needs_vision=needs_vision,
                token_estimate=token_estimate,
                allowed_models=allowed_models,
                exclude_models=exclude_models,
                language=language,
            )
            
            if result is not None:
                return RouteDecision(
                    provider=result.provider,
                    model=result.model,
                    score=result.score,
                    adjusted_score=result.adjusted_score,
                    cost_per_mtok=result.cost_per_mtok,
                    strategy=result.strategy,
                )
        
        # Fall back to simple routing (first matching model)
        return self._fallback_route(
            needs_tools=needs_tools,
            needs_vision=needs_vision,
            token_estimate=token_estimate,
            allowed_models=allowed_models,
            exclude_models=exclude_models,
            language=language,
        )
    
    def _fallback_route(
        self,
        needs_tools: bool = False,
        needs_vision: bool = False,
        token_estimate: int = 0,
        allowed_models: list[str] | None = None,
        exclude_models: set[str] | None = None,
        language: str | None = None,
    ) -> RouteDecision | None:
        """Fallback routing when no vectors available — picks first matching model."""
        from backend.language import is_language_supported
        exclude_models = exclude_models or set()
        
        for provider, model in self._config.all_models():
            model_id = f"{provider.id}/{model.id}"
            
            if model_id in exclude_models:
                continue
            if not is_language_supported(language, model.supported_languages):
                continue
            if needs_tools and not model.supports_tools:
                continue
            if needs_vision and not model.supports_vision:
                continue
            if token_estimate > 0 and token_estimate > model.max_context:
                continue
            if allowed_models and model_id not in allowed_models and model.id not in allowed_models:
                continue
            
            cost = (model.input_cost_per_mtok + model.output_cost_per_mtok) / 2
            return RouteDecision(
                provider=provider,
                model=model,
                score=0.0,
                adjusted_score=0.0,
                cost_per_mtok=cost,
                strategy="fallback",
            )
        
        return None
