"""TieredRouter wrapper that enforces correct embedding injection.

TieredRouter.route() internally uses OpenAI embeddings (512-dim), but the GNN
was trained on Qwen3 embeddings (256-dim). EmbeddedRouter requires an
EmbeddingProvider at construction and validates dimension match, making
mismatched embeddings impossible rather than merely discouraged.
"""

from __future__ import annotations

import logging
import random
import time
from typing import TYPE_CHECKING, Protocol

import numpy as np

if TYPE_CHECKING:
    from embergraphrouter.inference import TieredRouter
    from embergraphrouter.inference.tiered_router import (
        InferenceTier,
        TierConfig,
        TieredRoutingResult,
        TieredThresholds,
    )

logger = logging.getLogger(__name__)


class EmbeddingProvider(Protocol):
    """Protocol for query embedding providers."""

    @property
    def dimension(self) -> int:
        """Return the embedding dimension."""
        ...

    def embed(self, query: str) -> np.ndarray:
        """Embed a query string into a vector.

        Args:
            query: The text to embed.

        Returns:
            numpy array of shape (dimension,).
        """
        ...


class EmbeddedRouter:
    """TieredRouter wrapper that validates and injects embeddings.

    Dimension validation at construction prevents silent mismatches.
    """

    __slots__ = ("_router", "_embedder")

    def __init__(self, router: "TieredRouter", embedder: EmbeddingProvider) -> None:
        """Initialize with router and embedder, validating dimensions.

        Args:
            router: TieredRouter with GNN model loaded.
            embedder: EmbeddingProvider that produces correct-dimension embeddings.

        Raises:
            ValueError: If embedder.dimension != router.embed_dim.
        """
        if embedder.dimension != router.embed_dim:
            raise ValueError(
                f"Embedding dimension mismatch: embedder={embedder.dimension}, "
                f"router={router.embed_dim}. The embedder must match the GNN's "
                f"training embedding dimension."
            )

        self._router = router
        self._embedder = embedder

        logger.info(
            "EmbeddedRouter initialized: dim=%d, thresholds=%s",
            embedder.dimension,
            router.thresholds,
        )

    @property
    def thresholds(self) -> "TieredThresholds":
        """Expose thresholds for inspection (read-only)."""
        return self._router.thresholds

    @property
    def tier_configs(self) -> dict["InferenceTier", "TierConfig"]:
        """Expose tier configs for inspection (read-only)."""
        return self._router.tier_configs

    @property
    def embed_dim(self) -> int:
        """Return the embedding dimension."""
        return self._embedder.dimension

    def route(self, query: str) -> "TieredRoutingResult":
        """Route query using the injected embedder.

        Args:
            query: Query text to route.

        Returns:
            TieredRoutingResult with tier, signals, and selected model.
        """
        embed_start = time.perf_counter()
        embedding = self._embedder.embed(query)
        embed_latency = (time.perf_counter() - embed_start) * 1000

        return self.route_with_embedding(embedding, embed_latency)

    def route_with_embedding(
        self,
        embedding: np.ndarray,
        embed_latency_ms: float = 0.0,
    ) -> "TieredRoutingResult":
        """Route using a pre-computed embedding.

        This enables batch workflows: embed N queries in one GPU call via
        embed_batch(), then route each embedding individually through the GNN.

        Args:
            embedding: Pre-computed query embedding of shape (dimension,).
            embed_latency_ms: Time spent computing the embedding (for logging).

        Returns:
            TieredRoutingResult with tier, signals, and selected model.
        """
        from embergraphrouter.inference.tiered_router import (
            InferenceTier,
            TieredRoutingResult,
        )

        # Extract signals and classify
        classify_start = time.perf_counter()
        signals = self._router.extract_signals(embedding)
        tier = signals.classify_tier(self._router.thresholds)
        classify_latency = (time.perf_counter() - classify_start) * 1000

        config = self._router.tier_configs[tier]

        if config.use_ensemble:
            selected_model = f"eb1:{config.ensemble_plan}"
        elif tier == InferenceTier.DISCRIMINATIVE:
            selected_model = signals.top1_model
        elif tier == InferenceTier.OUT_OF_DISTRIBUTION:
            if (
                self._router.thresholds.ood_ensemble_probability > 0
                and random.random() < self._router.thresholds.ood_ensemble_probability
            ):
                uncertain_config = self._router.tier_configs.get(InferenceTier.UNCERTAIN)
                if uncertain_config and uncertain_config.use_ensemble:
                    selected_model = f"eb1:{uncertain_config.ensemble_plan}"
                else:
                    selected_model = config.model
            else:
                selected_model = config.model
        else:
            selected_model = config.model

        rationale = self._router._build_rationale(tier, signals)

        return TieredRoutingResult(
            tier=tier,
            signals=signals,
            selected_model=selected_model,
            selected_config=config,
            embedding_latency_ms=embed_latency_ms,
            classification_latency_ms=classify_latency,
            rationale=rationale,
        )
