"""EB1 Router - Production-ready inference with OpenAI embedding integration.

This module provides end-to-end routing with:
- OpenAI text-embedding-3-large integration (512-dim Matryoshka)
- Pre-cached heterogeneous graph structure
- Async embedding support for high throughput
- LRU caching for repeated queries
- Multiple routing strategies (argmax, cost-optimized, quality-optimized, utility)

Routing Strategies:
    1. route(): Pure argmax on P(correct) - simple baseline
    2. route_with_cost_optimization(): High entropy → cheap model (save cost)
    3. route_with_quality_optimization(): Low confidence → premium model (max quality)
    4. route_with_expected_utility(): Maximize weighted utility function

Example:
    from embergraphrouter.inference import EB1Router

    # Initialize once (loads model, caches graph structure)
    router = EB1Router.from_registry("eb1-gnn-preview-12-19")

    # === Strategy 1: Pure argmax (baseline) ===
    result = router.route("What is the capital of France?")
    print(f"Route to: {result.best_model} (confidence: {result.confidence:.2f})")

    # === Strategy 2: Cost optimization (DeepMind-style) ===
    # When uncertain, use cheap model to save cost
    result = router.route_with_cost_optimization(
        "Explain quantum computing",
        entropy_threshold=0.6,  # High entropy → uncertain
        cheap_model="GPT-5 Nano",
    )

    # === Strategy 3: Quality optimization (Anthropic-style) ===
    # When uncertain, escalate to premium model for safety
    result = router.route_with_quality_optimization(
        "Debug this critical production code",
        confidence_threshold=0.5,  # Low confidence → escalate
        premium_model="Claude Opus 4.5",
    )

    # === Strategy 4: Expected utility (OpenAI-style) ===
    # Optimize weighted combination of quality, cost, latency
    result = router.route_with_expected_utility(
        "Summarize this article",
        quality_weight=0.7,
        cost_weight=0.2,
        latency_weight=0.1,
    )

    # Async batch routing
    results = await router.route_batch_async([
        "Solve x^2 + 2x + 1 = 0",
        "Write a Python function to sort a list",
        "What is quantum entanglement?",
    ])
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import pickle
import random
import time
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import torch
import torch.nn.functional as F

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from numpy.typing import NDArray


@dataclass
class RoutingResult:
    """Result of a routing decision."""

    best_model: str
    best_model_idx: int
    confidence: float
    entropy: float  # Normalized [0, 1]
    scores: dict[str, float]
    probabilities: dict[str, float]
    embedding_latency_ms: float = 0.0
    inference_latency_ms: float = 0.0

    @property
    def total_latency_ms(self) -> float:
        return self.embedding_latency_ms + self.inference_latency_ms


@dataclass
class EB1RouterConfig:
    """Configuration for EB1 Router."""

    embed_dim: int = 512
    hidden_dim: int = 384
    num_gat_layers: int = 3
    num_gat_heads: int = 4
    dropout: float = 0.4
    edge_feature_dim: int = 6
    num_llms: int = 7

    # Inference preferences
    perf_weight: float = 0.7
    cost_weight: float = 0.2
    latency_weight: float = 0.1

    # OpenAI embedding config
    openai_model: str = "text-embedding-3-large"
    openai_dimensions: int = 512


class EmbeddingClient:
    """OpenAI embedding client with caching and exponential backoff.

    Implements retry with exponential backoff and jitter for rate limit handling.
    Follows Dean/Ghemawat principles: bounded retries, predictable behavior.
    """

    # Retry configuration (class-level for consistency)
    MAX_RETRIES = 5
    BASE_DELAY_SECONDS = 1.0
    MAX_DELAY_SECONDS = 60.0
    JITTER_FACTOR = 0.5  # ±50% jitter

    def __init__(
        self,
        model: str = "text-embedding-3-large",
        dimensions: int = 512,
        cache_size: int = 10000,
    ):
        self.model = model
        self.dimensions = dimensions
        self._cache: dict[str, np.ndarray] = {}
        self._cache_size = cache_size
        self._client = None
        self._async_client = None

    def _compute_backoff(self, attempt: int) -> float:
        """Compute backoff delay with exponential growth and jitter.

        Args:
            attempt: Zero-indexed retry attempt number.

        Returns:
            Delay in seconds with jitter applied.
        """
        # Exponential backoff: base * 2^attempt
        delay = self.BASE_DELAY_SECONDS * (2**attempt)
        delay = min(delay, self.MAX_DELAY_SECONDS)

        # Add jitter: delay * (1 ± jitter_factor/2)
        jitter = delay * self.JITTER_FACTOR * (random.random() - 0.5)
        return delay + jitter

    def _get_api_key(self) -> str:
        """Get OpenAI API key from Ember's unified configuration.

        Follows AGENTS.md: credentials from single canonical source (~/.ember/config.yaml).

        Raises:
            RuntimeError: If ember.core.credentials is unavailable or key retrieval fails.
        """
        from ember.core.credentials import get_api_key

        return get_api_key("openai")

    def _get_client(self):
        """Lazy-load OpenAI client with credentials from Ember config."""
        if self._client is None:
            try:
                from openai import OpenAI

                api_key = self._get_api_key()
                self._client = OpenAI(api_key=api_key)
            except ImportError as err:
                raise ImportError(
                    "openai package required for embedding. Install with: pip install openai"
                ) from err
        return self._client

    def _get_async_client(self):
        """Lazy-load async OpenAI client with credentials from Ember config."""
        if self._async_client is None:
            try:
                from openai import AsyncOpenAI

                api_key = self._get_api_key()
                self._async_client = AsyncOpenAI(api_key=api_key)
            except ImportError as err:
                raise ImportError(
                    "openai package required for embedding. Install with: pip install openai"
                ) from err
        return self._async_client

    def _cache_key(self, text: str) -> str:
        """Generate cache key for text."""
        return hashlib.md5(text.encode()).hexdigest()

    def _is_retryable_error(self, error: Exception) -> bool:
        """Check if an error is retryable (rate limit, transient failure)."""
        error_str = str(error).lower()
        # OpenAI rate limit errors
        if "rate" in error_str and "limit" in error_str:
            return True
        if "429" in error_str:  # HTTP 429 Too Many Requests
            return True
        if "503" in error_str or "502" in error_str:  # Service unavailable
            return True
        if "timeout" in error_str:
            return True
        return False

    def embed(self, text: str) -> tuple[np.ndarray, float]:
        """Get embedding for text, using cache if available.

        Implements retry with exponential backoff for rate limits.

        Returns:
            Tuple of (embedding array, latency in ms).

        Raises:
            RuntimeError: After MAX_RETRIES failures.
        """
        key = self._cache_key(text)
        if key in self._cache:
            return self._cache[key], 0.0

        start = time.perf_counter()
        client = self._get_client()
        last_error: Exception | None = None

        for attempt in range(self.MAX_RETRIES):
            try:
                response = client.embeddings.create(
                    model=self.model,
                    input=text,
                    dimensions=self.dimensions,
                )
                latency_ms = (time.perf_counter() - start) * 1000

                embedding = np.array(response.data[0].embedding, dtype=np.float32)

                # Cache with LRU eviction
                if len(self._cache) >= self._cache_size:
                    oldest_key = next(iter(self._cache))
                    del self._cache[oldest_key]
                self._cache[key] = embedding

                if attempt > 0:
                    logger.info(
                        "Embedding succeeded after %d retries (%.1fms)",
                        attempt,
                        latency_ms,
                    )

                return embedding, latency_ms

            except Exception as e:
                last_error = e
                if not self._is_retryable_error(e):
                    raise  # Non-retryable error, fail immediately

                if attempt < self.MAX_RETRIES - 1:
                    delay = self._compute_backoff(attempt)
                    logger.warning(
                        "Embedding request failed (attempt %d/%d): %s. Retrying in %.1fs...",
                        attempt + 1,
                        self.MAX_RETRIES,
                        str(e)[:100],
                        delay,
                    )
                    time.sleep(delay)

        # All retries exhausted
        raise RuntimeError(
            f"Embedding failed after {self.MAX_RETRIES} attempts: {last_error}"
        ) from last_error

    async def embed_async(self, text: str) -> tuple[np.ndarray, float]:
        """Async version of embed with retry and exponential backoff.

        Returns:
            Tuple of (embedding array, latency in ms).

        Raises:
            RuntimeError: After MAX_RETRIES failures.
        """
        key = self._cache_key(text)
        if key in self._cache:
            return self._cache[key], 0.0

        start = time.perf_counter()
        client = self._get_async_client()
        last_error: Exception | None = None

        for attempt in range(self.MAX_RETRIES):
            try:
                response = await client.embeddings.create(
                    model=self.model,
                    input=text,
                    dimensions=self.dimensions,
                )
                latency_ms = (time.perf_counter() - start) * 1000

                embedding = np.array(response.data[0].embedding, dtype=np.float32)

                if len(self._cache) >= self._cache_size:
                    oldest_key = next(iter(self._cache))
                    del self._cache[oldest_key]
                self._cache[key] = embedding

                if attempt > 0:
                    logger.info(
                        "Async embedding succeeded after %d retries (%.1fms)",
                        attempt,
                        latency_ms,
                    )

                return embedding, latency_ms

            except Exception as e:
                last_error = e
                if not self._is_retryable_error(e):
                    raise

                if attempt < self.MAX_RETRIES - 1:
                    delay = self._compute_backoff(attempt)
                    logger.warning(
                        "Async embedding failed (attempt %d/%d): %s. Retrying in %.1fs...",
                        attempt + 1,
                        self.MAX_RETRIES,
                        str(e)[:100],
                        delay,
                    )
                    await asyncio.sleep(delay)

        raise RuntimeError(
            f"Async embedding failed after {self.MAX_RETRIES} attempts: {last_error}"
        ) from last_error

    async def embed_batch_async(self, texts: list[str]) -> tuple[list[np.ndarray], float]:
        """Embed multiple texts with retry and exponential backoff.

        Batches uncached texts for efficiency while handling rate limits.

        Returns:
            Tuple of (list of embeddings, total latency in ms).

        Raises:
            RuntimeError: After MAX_RETRIES failures.
        """
        start = time.perf_counter()

        # Check cache first
        results: list[np.ndarray | None] = [None] * len(texts)
        uncached_indices = []
        uncached_texts = []

        for i, text in enumerate(texts):
            key = self._cache_key(text)
            if key in self._cache:
                results[i] = self._cache[key]
            else:
                uncached_indices.append(i)
                uncached_texts.append(text)

        # Batch embed uncached texts with retry
        if uncached_texts:
            client = self._get_async_client()
            last_error: Exception | None = None

            for attempt in range(self.MAX_RETRIES):
                try:
                    response = await client.embeddings.create(
                        model=self.model,
                        input=uncached_texts,
                        dimensions=self.dimensions,
                    )

                    for idx, (orig_idx, text) in enumerate(
                        zip(uncached_indices, uncached_texts, strict=False)
                    ):
                        embedding = np.array(response.data[idx].embedding, dtype=np.float32)
                        results[orig_idx] = embedding

                        # Cache
                        key = self._cache_key(text)
                        if len(self._cache) < self._cache_size:
                            self._cache[key] = embedding

                    if attempt > 0:
                        logger.info(
                            "Batch embedding succeeded after %d retries",
                            attempt,
                        )
                    break  # Success, exit retry loop

                except Exception as e:
                    last_error = e
                    if not self._is_retryable_error(e):
                        raise

                    if attempt < self.MAX_RETRIES - 1:
                        delay = self._compute_backoff(attempt)
                        logger.warning(
                            "Batch embedding failed (attempt %d/%d, %d texts): %s. "
                            "Retrying in %.1fs...",
                            attempt + 1,
                            self.MAX_RETRIES,
                            len(uncached_texts),
                            str(e)[:100],
                            delay,
                        )
                        await asyncio.sleep(delay)
            else:
                # All retries exhausted (loop completed without break)
                raise RuntimeError(
                    f"Batch embedding failed after {self.MAX_RETRIES} attempts: {last_error}"
                ) from last_error

        latency_ms = (time.perf_counter() - start) * 1000
        return results, latency_ms  # type: ignore[return-value]


class EB1Router:
    """Production-ready EB1 Router with end-to-end inference."""

    def __init__(
        self,
        model: torch.nn.Module,
        llm_embeddings: torch.Tensor,
        llm_names: list[str],
        llm_metadata: dict[str, dict],
        config: EB1RouterConfig,
        device: str = "cpu",
    ):
        self.model = model
        self.llm_embeddings = llm_embeddings.to(device)
        self.llm_names = llm_names
        self.llm_metadata = llm_metadata
        self.config = config
        self.device = device
        self.num_llms = len(llm_names)

        # Initialize embedding client
        self._embedding_client = EmbeddingClient(
            model=config.openai_model,
            dimensions=config.openai_dimensions,
        )

        # Pre-build static graph components
        self._edge_index = self._build_edge_index()
        self._edge_features = self._build_edge_features()

        # Set model to eval mode
        self.model.eval()

    def _build_edge_index(self) -> torch.Tensor:
        """Build static edge index for heterogeneous graph.

        Graph structure matches training:
        - Nodes 0..num_queries-1: Query nodes
        - Nodes num_queries..num_queries+num_llms-1: LLM nodes
        - Edges: Each query connects to all LLMs

        For single query (num_queries=1):
        - Node 0: Query
        - Nodes 1..7: LLMs
        - source_nodes = [0, 0, 0, 0, 0, 0, 0]
        - dest_node_ids = [1, 2, 3, 4, 5, 6, 7]
        """
        num_queries = 1  # Single query inference
        # source_nodes: query index repeated for each LLM
        source_nodes = [0] * self.num_llms
        # dest_nodes: LLM indices (will be offset by num_queries)
        dest_nodes = list(range(self.num_llms))
        # Apply offset: dest_node_ids = i + 1 + source_nodes[-1] = i + 1
        dest_node_ids = [i + num_queries for i in dest_nodes]

        edge_index = torch.tensor(
            [source_nodes, dest_node_ids],
            dtype=torch.long,
            device=self.device,
        )
        return edge_index

    def _build_edge_features(self) -> torch.Tensor:
        """Build edge features with preference weights.

        Edge features: [perf, cost, latency, pref_perf, pref_cost, pref_latency]
        """
        num_edges = self.num_llms  # One edge per LLM (single query)
        edge_features = torch.zeros(num_edges, 6, device=self.device)

        # Set preference weights (same for all edges)
        edge_features[:, 3] = self.config.perf_weight
        edge_features[:, 4] = self.config.cost_weight
        edge_features[:, 5] = self.config.latency_weight

        # Performance features (quality = 1.0 for now, actual values from metadata)
        edge_features[:, 0] = 1.0

        return edge_features

    @classmethod
    def from_registry(
        cls,
        registry_name: str,
        registry_root: str | Path | None = None,
        device: str = "cpu",
    ) -> EB1Router:
        """Load router from registry.

        Args:
            registry_name: Name of registry entry (e.g., "eb1-gnn-preview-12-19").
            registry_root: Root directory. Defaults to embergraphrouter/registry/.
            device: Device for inference ("cpu" or "cuda").

        Returns:
            Initialized EB1Router.
        """
        if registry_root is None:
            registry_root = Path(__file__).parent.parent / "registry"
        registry_root = Path(registry_root)

        registry_path = registry_root / registry_name
        if not registry_path.exists():
            raise FileNotFoundError(f"Registry not found: {registry_path}")

        # Load config
        import yaml

        config_path = registry_path / "config.yaml"
        with open(config_path) as f:
            config_dict = yaml.safe_load(f)

        config = EB1RouterConfig(
            embed_dim=config_dict.get("embed_dim", 512),
            hidden_dim=config_dict.get("hidden_dim", 384),
            num_gat_layers=config_dict.get("num_gat_layers", 3),
            num_gat_heads=config_dict.get("num_gat_heads", 4),
            dropout=config_dict.get("dropout", 0.4),
            edge_feature_dim=config_dict.get("edge_feature_dim", 6),
            num_llms=config_dict.get("llm_num", 7),
        )

        # Load LLM embeddings
        emb_path = registry_path / "llm_embeddings.pkl"
        with open(emb_path, "rb") as f:
            llm_embeddings_np = pickle.load(f)
        llm_embeddings = torch.from_numpy(llm_embeddings_np).float()

        # Load LLM descriptions
        desc_path = registry_path / "llm_descriptions.json"
        with open(desc_path) as f:
            llm_descriptions = json.load(f)

        llm_names = list(llm_descriptions.keys())
        llm_metadata = {
            name: {
                "cost": desc.get("cost", 0.0),
                "latency": desc.get("latency", 0.0),
            }
            for name, desc in llm_descriptions.items()
        }

        # Initialize model architecture
        # Handle both installed package and local development
        try:
            from embergraphrouter.model.graph_nn import GATEncoderDecoderNet
        except ModuleNotFoundError:
            from model.graph_nn import GATEncoderDecoderNet

        # Load weights first to detect architecture from state_dict
        model_path = registry_path / "model.pth"
        state_dict = torch.load(model_path, map_location=device, weights_only=True)
        use_score_head = "score_head.weight" in state_dict

        model = GATEncoderDecoderNet(
            query_feature_dim=config.embed_dim,
            llm_feature_dim=config.embed_dim,
            hidden_features=config.hidden_dim,
            in_edges=config.edge_feature_dim,
            num_gnn_layers=config.num_gat_layers,
            num_attention_heads=config.num_gat_heads,
            dropout=config.dropout,
            hidden_multiplier=2,  # Match training config
            use_contrastive=True,  # Checkpoint has contrastive_encoder
            use_score_head=use_score_head,
            use_difficulty_head=True,  # Checkpoint has difficulty_head
            use_uncertainty_head=False,  # Checkpoint does NOT have uncertainty_proj
        )

        # Load weights (strict=False to ignore deprecated difficulty_head keys)
        model.load_state_dict(state_dict, strict=False)
        model.to(device)

        return cls(
            model=model,
            llm_embeddings=llm_embeddings,
            llm_names=llm_names,
            llm_metadata=llm_metadata,
            config=config,
            device=device,
        )

    def route(self, query: str) -> RoutingResult:
        """Route a query to the best LLM.

        This is the main entry point for routing. It:
        1. Embeds the query using OpenAI API
        2. Runs GNN inference
        3. Returns routing decision with confidence

        Args:
            query: The query text to route.

        Returns:
            RoutingResult with best model and confidence.
        """
        # Get embedding
        embedding, embed_latency = self._embedding_client.embed(query)

        # Run inference
        start = time.perf_counter()
        result = self._route_from_embedding(embedding)
        inference_latency = (time.perf_counter() - start) * 1000

        result.embedding_latency_ms = embed_latency
        result.inference_latency_ms = inference_latency

        return result

    async def route_async(self, query: str) -> RoutingResult:
        """Async version of route."""
        embedding, embed_latency = await self._embedding_client.embed_async(query)

        start = time.perf_counter()
        result = self._route_from_embedding(embedding)
        inference_latency = (time.perf_counter() - start) * 1000

        result.embedding_latency_ms = embed_latency
        result.inference_latency_ms = inference_latency

        return result

    async def route_batch_async(self, queries: list[str]) -> list[RoutingResult]:
        """Route multiple queries concurrently.

        Batches embedding API calls for efficiency.
        """
        embeddings, embed_latency = await self._embedding_client.embed_batch_async(queries)
        per_query_embed_latency = embed_latency / len(queries)

        results = []
        for embedding in embeddings:
            start = time.perf_counter()
            result = self._route_from_embedding(embedding)
            inference_latency = (time.perf_counter() - start) * 1000

            result.embedding_latency_ms = per_query_embed_latency
            result.inference_latency_ms = inference_latency
            results.append(result)

        return results

    def route_from_embedding(self, embedding: NDArray[np.float32] | torch.Tensor) -> RoutingResult:
        """Route from pre-computed embedding (skips OpenAI API call)."""
        return self._route_from_embedding(embedding)

    def _route_from_embedding(self, embedding: NDArray[np.float32] | torch.Tensor) -> RoutingResult:
        """Internal routing from embedding."""
        if isinstance(embedding, np.ndarray):
            embedding = torch.from_numpy(embedding).float()
        embedding = embedding.to(self.device)

        if embedding.dim() == 1:
            embedding = embedding.unsqueeze(0)

        with torch.no_grad():
            logits = self._forward(embedding)

        probs = F.softmax(logits, dim=-1).squeeze(0).cpu().numpy()
        logits_np = logits.squeeze(0).cpu().numpy()

        # Compute normalized entropy
        entropy = -np.sum(probs * np.log(probs + 1e-8))
        max_entropy = np.log(self.num_llms)
        normalized_entropy = entropy / max_entropy

        best_idx = int(np.argmax(probs))
        best_model = self.llm_names[best_idx]
        confidence = float(probs[best_idx])

        scores = {name: float(logits_np[i]) for i, name in enumerate(self.llm_names)}
        probabilities = {name: float(probs[i]) for i, name in enumerate(self.llm_names)}

        return RoutingResult(
            best_model=best_model,
            best_model_idx=best_idx,
            confidence=confidence,
            entropy=normalized_entropy,
            scores=scores,
            probabilities=probabilities,
        )

    def _forward(self, query_embedding: torch.Tensor) -> torch.Tensor:
        """Run forward pass for single-query inference.

        For single-query inference, GAT message passing collapses all LLM
        representations to identical values (because all edges have the same
        source). Instead, we compute predictions directly from aligned
        embeddings (x_initial) without GNN processing.

        This matches the model's prediction mechanism but skips the GNN
        layers that require batch context to be meaningful.
        """
        task_embedding = query_embedding  # (1, embed_dim)

        # Get aligned embeddings (before GNN)
        x_initial = self.model.model_align(task_embedding, query_embedding, self.llm_embeddings)
        # x_initial: (1 + num_llms, hidden_dim) = (8, 768)
        # - Node 0: query node
        # - Nodes 1-7: LLM nodes

        # For single query, edge_index_predict is:
        # - source: [0, 0, 0, 0, 0, 0, 0] (query)
        # - dest: [1, 2, 3, 4, 5, 6, 7] (LLMs)
        edge_index_predict = self._edge_index

        # Get representations for edge prediction
        # query_repr: same query repeated for each edge
        # llm_repr: different LLM for each edge
        query_repr = x_initial[edge_index_predict[0]]  # (7, 768)
        llm_repr = x_initial[edge_index_predict[1]]  # (7, 768) - differentiated!

        # Edge prediction via element-wise product + scoring
        edge_repr = query_repr * llm_repr  # (num_llms, channels)
        if self.model.score_head is not None:
            logits = self.model.score_head(edge_repr).squeeze(-1)
        else:
            # Match training's .mean(), not .sum() (fixes 512x magnitude mismatch)
            logits = edge_repr.mean(dim=-1)

        # Reshape to (1, num_llms) for consistency
        logits = logits.unsqueeze(0)

        return logits

    def route_with_cost_optimization(
        self,
        query: str,
        entropy_threshold: float = 0.6,
        cheap_model: str | None = None,
    ) -> RoutingResult:
        """Route with entropy-based cost optimization.

        When prediction entropy is high (models are similar), use cheap model.
        When entropy is low (router found a winner), trust prediction.

        Args:
            query: Query text.
            entropy_threshold: Threshold for using cheap model.
            cheap_model: Fallback model name. Defaults to cheapest.

        Returns:
            RoutingResult with potentially overridden best_model.
        """
        result = self.route(query)

        if result.entropy > entropy_threshold:
            if cheap_model is None:
                # Find cheapest model
                cheap_model = min(
                    self.llm_metadata.items(),
                    key=lambda x: x[1].get("cost", float("inf")),
                )[0]

            if cheap_model in self.llm_names:
                return RoutingResult(
                    best_model=cheap_model,
                    best_model_idx=self.llm_names.index(cheap_model),
                    confidence=result.probabilities.get(cheap_model, 0.0),
                    entropy=result.entropy,
                    scores=result.scores,
                    probabilities=result.probabilities,
                    embedding_latency_ms=result.embedding_latency_ms,
                    inference_latency_ms=result.inference_latency_ms,
                )

        return result

    def route_with_quality_optimization(
        self,
        query: str,
        confidence_threshold: float = 0.5,
        premium_model: str | None = None,
    ) -> RoutingResult:
        """Route with confidence-gated escalation to premium model.

        When prediction confidence is low (router is uncertain), escalate to
        premium model for higher quality. When confident, trust the prediction.

        This implements the Anthropic "ask the expert when uncertain" pattern:
        - Low confidence: Router can't distinguish → use most capable model
        - High confidence: Router found a winner → trust prediction

        Args:
            query: Query text.
            confidence_threshold: Minimum confidence to trust prediction.
                Below this, escalate to premium. Default 0.5.
            premium_model: Fallback model name. Defaults to most expensive.

        Returns:
            RoutingResult with potentially escalated best_model.
        """
        result = self.route(query)

        if result.confidence < confidence_threshold:
            if premium_model is None:
                # Find most expensive model (proxy for most capable)
                premium_model = max(
                    self.llm_metadata.items(),
                    key=lambda x: x[1].get("cost", 0.0),
                )[0]

            if premium_model in self.llm_names:
                return RoutingResult(
                    best_model=premium_model,
                    best_model_idx=self.llm_names.index(premium_model),
                    confidence=result.probabilities.get(premium_model, 0.0),
                    entropy=result.entropy,
                    scores=result.scores,
                    probabilities=result.probabilities,
                    embedding_latency_ms=result.embedding_latency_ms,
                    inference_latency_ms=result.inference_latency_ms,
                )

        return result

    def route_with_expected_utility(
        self,
        query: str,
        quality_weight: float = 0.7,
        cost_weight: float = 0.2,
        latency_weight: float = 0.1,
    ) -> RoutingResult:
        """Route by maximizing expected utility across quality, cost, and latency.

        Instead of pure argmax on P(correct), this computes:
            utility[i] = quality_weight * P(correct|LLM_i)
                       - cost_weight * normalized_cost[i]
                       - latency_weight * normalized_latency[i]

        This implements the OpenAI cost-quality tradeoff optimization:
        - High quality_weight: Prioritize correctness
        - High cost_weight: Prefer cheaper models
        - High latency_weight: Prefer faster models

        Args:
            query: Query text.
            quality_weight: Weight for P(correct). Default 0.7.
            cost_weight: Weight for cost penalty. Default 0.2.
            latency_weight: Weight for latency penalty. Default 0.1.

        Returns:
            RoutingResult with utility-maximizing best_model.
        """
        result = self.route(query)

        # Gather costs and latencies
        costs = np.array(
            [self.llm_metadata.get(name, {}).get("cost", 0.0) for name in self.llm_names]
        )
        latencies = np.array(
            [self.llm_metadata.get(name, {}).get("latency", 0.0) for name in self.llm_names]
        )

        # Normalize to [0, 1] range
        cost_range = costs.max() - costs.min()
        if cost_range > 0:
            normalized_costs = (costs - costs.min()) / cost_range
        else:
            normalized_costs = np.zeros_like(costs)

        latency_range = latencies.max() - latencies.min()
        if latency_range > 0:
            normalized_latencies = (latencies - latencies.min()) / latency_range
        else:
            normalized_latencies = np.zeros_like(latencies)

        # Compute utility for each LLM
        probs = np.array([result.probabilities[name] for name in self.llm_names])
        utility = (
            quality_weight * probs
            - cost_weight * normalized_costs
            - latency_weight * normalized_latencies
        )

        best_idx = int(np.argmax(utility))
        best_model = self.llm_names[best_idx]

        return RoutingResult(
            best_model=best_model,
            best_model_idx=best_idx,
            confidence=result.probabilities[best_model],
            entropy=result.entropy,
            scores=result.scores,
            probabilities=result.probabilities,
            embedding_latency_ms=result.embedding_latency_ms,
            inference_latency_ms=result.inference_latency_ms,
        )

    def benchmark(self, num_queries: int = 100) -> dict:
        """Benchmark inference latency.

        Args:
            num_queries: Number of queries to benchmark.

        Returns:
            Dict with latency statistics.
        """
        # Generate random embeddings (skip API for benchmarking)
        embeddings = [
            np.random.randn(self.config.embed_dim).astype(np.float32) for _ in range(num_queries)
        ]

        latencies = []
        for emb in embeddings:
            start = time.perf_counter()
            _ = self._route_from_embedding(emb)
            latencies.append((time.perf_counter() - start) * 1000)

        latencies = np.array(latencies)

        return {
            "num_queries": num_queries,
            "mean_ms": float(np.mean(latencies)),
            "std_ms": float(np.std(latencies)),
            "p50_ms": float(np.percentile(latencies, 50)),
            "p95_ms": float(np.percentile(latencies, 95)),
            "p99_ms": float(np.percentile(latencies, 99)),
            "throughput_qps": 1000 / np.mean(latencies),
        }
