"""Vector-based router — routes prompts to models using learned embeddings."""

from __future__ import annotations

import json
import logging
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import numpy as np

from backend.config import EcateConfig, ModelConfig, ProviderConfig
from backend.router.embedder import get_embedder

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


@dataclass
class RouteDecision:
    """Result of routing decision."""
    provider: ProviderConfig
    model: ModelConfig
    score: float  # Raw dot product score
    adjusted_score: float  # Score after strategy adjustment
    cost_per_mtok: float
    strategy: str
    all_scores: dict[str, float]  # model_id -> score for debugging


class VectorRouter:
    """Routes requests to optimal model using learned embeddings.
    
    Each model has a learned 512d vector. Routing works by:
    1. Embed the prompt locally (fast, ~10ms)
    2. Dot product with each model's vector → scores
    3. Apply strategy (auto/best/cheap) to select model
    
    Strategies:
    - "best": Pick highest scoring model regardless of cost
    - "cheap": Pick cheapest model above threshold
    - "auto": Balance score vs cost (score - λ*log(cost))
    """
    
    def __init__(
        self, 
        config: EcateConfig,
        vectors_path: Optional[Path] = None,
        cost_weight: float = 0.3,
        score_threshold: float = 0.0,
    ):
        self._config = config
        self._vectors_path = vectors_path or Path("data/model_vectors.json")
        self._cost_weight = cost_weight
        self._score_threshold = score_threshold
        
        # model_id -> 512d vector
        self._model_vectors: dict[str, np.ndarray] = {}
        # model_id -> cost per M tokens (avg of in/out)
        self._model_costs: dict[str, float] = {}
        
        self._embedder = get_embedder()
        self._load_vectors()
        self._load_costs()
    
    def _load_vectors(self) -> None:
        """Load model vectors from JSON file."""
        if not self._vectors_path.exists():
            logger.warning(f"No model vectors at {self._vectors_path}, routing will use fallback")
            return
        
        try:
            with open(self._vectors_path) as f:
                data = json.load(f)
            
            expected_dim = data.get("embed_dim", 512)
            loaded = 0
            skipped = 0
            
            for model_id, vec in data.get("vectors", {}).items():
                arr = np.array(vec, dtype=np.float32)
                if arr.shape[0] != expected_dim:
                    logger.warning(f"Vector dim mismatch for {model_id}: {arr.shape[0]} vs {expected_dim}, skipping")
                    skipped += 1
                    continue
                self._model_vectors[model_id] = arr
                loaded += 1
            
            if skipped:
                logger.warning(f"Skipped {skipped} vectors due to dimension mismatch")
            logger.info(f"Loaded {loaded} model vectors from {self._vectors_path}")
        except Exception as e:
            logger.error(f"Failed to load model vectors: {e}")
    
    def reload_vectors(self) -> None:
        """Reload vectors from disk (call after calibration updates)."""
        self._model_vectors.clear()
        self._load_vectors()
        self._load_costs()
    
    def _load_costs(self) -> None:
        """Load cost per model from config."""
        for provider, model in self._config.all_models():
            model_id = f"{provider.id}/{model.id}"
            # Average of input and output cost
            cost = (model.input_cost_per_mtok + model.output_cost_per_mtok) / 2
            self._model_costs[model_id] = cost
    
    def _get_vector(self, model_id: str) -> Optional[np.ndarray]:
        """Get model vector, trying full id then just model name."""
        if model_id in self._model_vectors:
            return self._model_vectors[model_id]
        # Try without provider prefix
        model_name = model_id.split("/")[-1] if "/" in model_id else model_id
        for k, v in self._model_vectors.items():
            if k.endswith(f"/{model_name}") or k == model_name:
                return v
        return None
    
    def has_vectors(self) -> bool:
        """Check if any model vectors are loaded."""
        return len(self._model_vectors) > 0
    
    def get_calibrated_models(self) -> list[str]:
        """Get list of model IDs that have learned vectors."""
        return list(self._model_vectors.keys())
    
    def route(
        self,
        prompt: str,
        strategy: str = "auto",
        needs_tools: bool = False,
        needs_vision: bool = False,
        token_estimate: int = 0,
        allowed_models: Optional[list[str]] = None,
        exclude_models: Optional[set[str]] = None,
        language: Optional[str] = None,
    ) -> Optional[RouteDecision]:
        """Route prompt to optimal model.
        
        Args:
            prompt: User prompt to route
            strategy: Routing strategy (auto/best/cheap)
            needs_tools: Filter to models supporting tools
            needs_vision: Filter to models supporting vision
            token_estimate: Estimated total tokens (for context limit)
            allowed_models: Whitelist of model IDs
            exclude_models: Blacklist of model IDs
            language: Detected language code or None (all models considered)
        
        Returns:
            RouteDecision or None if no suitable model
        """
        from backend.language import is_language_supported
        exclude_models = exclude_models or set()
        
        # Collect candidate models
        candidates: list[tuple[ProviderConfig, ModelConfig, str]] = []
        for provider, model in self._config.all_models():
            model_id = f"{provider.id}/{model.id}"
            
            # Exclusion filter
            if model_id in exclude_models:
                continue
            
            # Language filter
            if not is_language_supported(language, model.supported_languages):
                continue
            
            # Capability filters
            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
            
            candidates.append((provider, model, model_id))
        
        if not candidates:
            return None
        
        # Embed the prompt
        prompt_vec = self._embedder.embed(prompt)
        
        # Score each candidate
        scores: dict[str, float] = {}
        adjusted: dict[str, float] = {}  # Initialize for all strategies
        for provider, model, model_id in candidates:
            model_vec = self._get_vector(model_id)
            if model_vec is not None:
                # Dot product (both normalized → cosine similarity)
                score = float(np.dot(prompt_vec, model_vec))
            else:
                # No vector → neutral score
                score = 0.0
            scores[model_id] = score
        
        # Apply strategy
        if strategy == "best":
            # Pick highest score
            best_id = max(scores, key=lambda k: scores[k])
        
        elif strategy == "cheap":
            # Pick cheapest above threshold
            above_threshold = {k: v for k, v in scores.items() if v >= self._score_threshold}
            if not above_threshold:
                # Relax: pick cheapest overall
                above_threshold = scores
            best_id = min(above_threshold, key=lambda k: self._model_costs.get(k, 0))
        
        else:  # auto
            # Balance score vs cost: adjusted = score - λ * log(cost)
            for model_id, score in scores.items():
                cost = self._model_costs.get(model_id, 0.001)
                # log(cost) gives diminishing penalty as cost increases
                # +1 to avoid log(0), shift so free models get bonus
                adj = score - self._cost_weight * math.log(cost + 0.01)
                adjusted[model_id] = adj
            
            best_id = max(adjusted, key=lambda k: adjusted[k])
        
        # Find winner
        for provider, model, model_id in candidates:
            if model_id == best_id:
                return RouteDecision(
                    provider=provider,
                    model=model,
                    score=scores[model_id],
                    adjusted_score=adjusted.get(model_id, scores[model_id]) if strategy == "auto" else scores[model_id],
                    cost_per_mtok=self._model_costs.get(model_id, 0),
                    strategy=strategy,
                    all_scores=scores,
                )
        
        return None
    
    def add_model_vector(self, model_id: str, vector: np.ndarray) -> None:
        """Add or update a model's vector."""
        self._model_vectors[model_id] = np.array(vector, dtype=np.float32)
        self._save_vectors()
    
    def remove_model_vector(self, model_id: str) -> bool:
        """Remove a model's vector."""
        if model_id in self._model_vectors:
            del self._model_vectors[model_id]
            self._save_vectors()
            return True
        return False
    
    def _save_vectors(self) -> None:
        """Save model vectors to JSON file."""
        self._vectors_path.parent.mkdir(parents=True, exist_ok=True)
        data = {
            "version": 1,
            "embed_dim": self._embedder.dim,
            "embedder_model": self._embedder._model_name,
            "vectors": {k: v.tolist() for k, v in self._model_vectors.items()},
        }
        with open(self._vectors_path, "w") as f:
            json.dump(data, f, indent=2)
        logger.info(f"Saved {len(self._model_vectors)} model vectors to {self._vectors_path}")
