"""kNN router -- Tier 1 embedding-based routing.

Fast classification for ~75% of requests.  Embeds the incoming query,
searches *K* nearest neighbours in the FAISS store, scores routes based on
multiple weighted factors, and either returns a confident decision or
escalates to Tier 2 (ModernBERT).

A 5% exploration rate deliberately routes requests randomly to gather
diverse training data.

Requires ``[smart]`` tier: ``faiss-cpu``, ``numpy``, ``onnxruntime``.
"""

from __future__ import annotations

import logging
import random
import time
from typing import TYPE_CHECKING

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Guarded imports
# ---------------------------------------------------------------------------

KNN_AVAILABLE: bool = False

try:
    import numpy as np

    from llmhosts.embeddings.pipeline import EMBEDDINGS_AVAILABLE
    from llmhosts.embeddings.store import FAISS_AVAILABLE

    if EMBEDDINGS_AVAILABLE and FAISS_AVAILABLE:
        KNN_AVAILABLE = True
except ImportError:
    np = None  # type: ignore[assignment]

if TYPE_CHECKING:
    import numpy as np

    from llmhosts.embeddings.pipeline import EmbeddingPipeline
    from llmhosts.embeddings.store import EmbeddingMetadata, EmbeddingStore, SearchResult
    from llmhosts.proxy.models import UnifiedRequest


# ============================================================================
# Pydantic result models
# ============================================================================


class ScoredRoute(BaseModel):
    """A candidate route with a composite score and per-factor breakdown."""

    model: str
    backend: str
    score: float  # 0.0-1.0 composite score
    factor_scores: dict[str, float] = Field(default_factory=dict)


class KNNRoutingResult(BaseModel):
    """The output of the kNN router tier."""

    action: str  # "route" | "escalate" | "explore"
    confidence: float  # 0.0-1.0
    suggested_model: str | None = None
    suggested_backend: str | None = None
    reasoning: str
    neighbor_count: int = 0
    exploration: bool = False
    scored_routes: list[ScoredRoute] = Field(default_factory=list)
    latency_ms: float = 0.0


# ============================================================================
# Factor weights
# ============================================================================

# Weights for the multi-factor scoring function.  Sum = 1.0.
_FACTOR_WEIGHTS: dict[str, float] = {
    "model_fit": 0.30,  # How well does this model handle similar queries?
    "latency": 0.20,  # Average latency for this route
    "cost": 0.20,  # Cost efficiency (lower = better)
    "quality": 0.20,  # User feedback scores
    "recency": 0.10,  # Prefer recent routing decisions
}


# ============================================================================
# KNNRouter
# ============================================================================


class KNNRouter:
    """Tier 1: kNN embedding router.

    Fast classification for ~75% of requests.
    Embed query -> K nearest lookup -> multi-factor score -> decide.

    If confidence >= ``CONFIDENCE_THRESHOLD``, route directly.
    If confidence < threshold, escalate to Tier 2 (ModernBERT).
    5% of requests are deliberate exploration (random route).

    Usage::

        knn = KNNRouter(pipeline, store)
        result = await knn.route(request)
        if result.action == "route":
            # use result.suggested_model / suggested_backend
        elif result.action == "escalate":
            # pass to Tier 2
    """

    CONFIDENCE_THRESHOLD: float = 0.85
    EXPLORATION_RATE: float = 0.05
    K_NEIGHBORS: int = 10

    # Minimum number of stored embeddings before kNN can make decisions.
    # Below this we always escalate to let other tiers decide while we
    # accumulate training data.
    MIN_STORE_SIZE: int = 20

    def __init__(self, pipeline: EmbeddingPipeline, store: EmbeddingStore) -> None:
        if not KNN_AVAILABLE:
            raise RuntimeError(
                "KNNRouter requires faiss-cpu, numpy, and onnxruntime. Install them with: pip install llmhosts[smart]"
            )
        self._pipeline = pipeline
        self._store = store

    # ------------------------------------------------------------------
    # Main entry point
    # ------------------------------------------------------------------

    async def route(self, request: UnifiedRequest) -> KNNRoutingResult:
        """Make a routing decision using kNN similarity search.

        Steps:
        1. Convert request messages to a single text.
        2. Embed the text.
        3. Check exploration trigger (5%).
        4. Search K nearest neighbours.
        5. Score candidate routes by multiple factors.
        6. If top score >= threshold, return route; else escalate.
        """
        start = time.perf_counter()

        # Pre-check: not enough data yet
        if self._store.size < self.MIN_STORE_SIZE:
            elapsed = (time.perf_counter() - start) * 1000
            return KNNRoutingResult(
                action="escalate",
                confidence=0.0,
                reasoning=(
                    f"Insufficient training data ({self._store.size}/{self.MIN_STORE_SIZE} embeddings). "
                    "Escalating to Tier 2 while accumulating data."
                ),
                neighbor_count=0,
                latency_ms=elapsed,
            )

        # Step 1 & 2: text -> embedding
        text = self._request_to_text(request)
        embedding = self._pipeline.embed(text)

        # Step 3: exploration check (5% of requests)
        if random.random() < self.EXPLORATION_RATE:
            elapsed = (time.perf_counter() - start) * 1000
            return KNNRoutingResult(
                action="explore",
                confidence=0.0,
                reasoning="Deliberate exploration (5% sampling) for training data diversity.",
                neighbor_count=0,
                exploration=True,
                latency_ms=elapsed,
            )

        # Step 4: kNN search
        neighbors = await self._store.search(embedding, k=self.K_NEIGHBORS)

        if not neighbors:
            elapsed = (time.perf_counter() - start) * 1000
            return KNNRoutingResult(
                action="escalate",
                confidence=0.0,
                reasoning="No neighbours found in store. Escalating to Tier 2.",
                neighbor_count=0,
                latency_ms=elapsed,
            )

        # Step 5: score candidate routes
        scored_routes = self._score_neighbors(neighbors, request)

        if not scored_routes:
            elapsed = (time.perf_counter() - start) * 1000
            return KNNRoutingResult(
                action="escalate",
                confidence=0.0,
                reasoning="No scorable routes from neighbours. Escalating to Tier 2.",
                neighbor_count=len(neighbors),
                latency_ms=elapsed,
            )

        # Step 6: decide
        best = scored_routes[0]
        elapsed = (time.perf_counter() - start) * 1000

        if best.score >= self.CONFIDENCE_THRESHOLD:
            return KNNRoutingResult(
                action="route",
                confidence=best.score,
                suggested_model=best.model,
                suggested_backend=best.backend,
                reasoning=(
                    f"kNN confident ({best.score:.2f} >= {self.CONFIDENCE_THRESHOLD}): "
                    f"route to {best.model} on {best.backend}. "
                    f"Factors: {self._format_factors(best.factor_scores)}"
                ),
                neighbor_count=len(neighbors),
                scored_routes=scored_routes,
                latency_ms=elapsed,
            )

        return KNNRoutingResult(
            action="escalate",
            confidence=best.score,
            suggested_model=best.model,
            suggested_backend=best.backend,
            reasoning=(
                f"kNN uncertain ({best.score:.2f} < {self.CONFIDENCE_THRESHOLD}): "
                f"best candidate is {best.model} on {best.backend}, "
                f"but escalating to Tier 2 for confirmation."
            ),
            neighbor_count=len(neighbors),
            scored_routes=scored_routes,
            latency_ms=elapsed,
        )

    # ------------------------------------------------------------------
    # Text extraction
    # ------------------------------------------------------------------

    def _request_to_text(self, request: UnifiedRequest) -> str:
        """Convert request messages to a single text for embedding.

        Strategy: concatenate the last few messages (system + user context)
        to capture the query semantics without overflowing the max sequence
        length.  Prioritise the most recent user message.
        """
        parts: list[str] = []
        # Include model name as a routing signal
        parts.append(f"[model:{request.model}]")

        # Gather message contents (last 5 messages max)
        messages = request.messages[-5:] if len(request.messages) > 5 else request.messages
        for msg in messages:
            content = msg.content
            if content is None:
                continue
            if isinstance(content, list):
                # Multi-part content (e.g. vision) -- extract text parts
                text_parts = [
                    block.get("text", "")
                    for block in content
                    if isinstance(block, dict) and block.get("type") == "text"
                ]
                content = " ".join(text_parts)
            if content:
                parts.append(f"[{msg.role}] {content}")

        return " ".join(parts)

    # ------------------------------------------------------------------
    # Multi-factor scoring
    # ------------------------------------------------------------------

    def _score_neighbors(self, neighbors: list[SearchResult], request: UnifiedRequest) -> list[ScoredRoute]:
        """Score candidate routes based on neighbour data.

        Groups neighbours by ``(model, backend)`` and computes a weighted
        composite score from five factors:

        - **model_fit** (0.30): similarity score * frequency of this route
          among neighbours -- higher means this model consistently handles
          similar queries.
        - **latency** (0.20): normalised inverse of average latency -- lower
          latency scores higher.
        - **cost** (0.20): normalised inverse of average cost -- cheaper is
          better.
        - **quality** (0.20): average user feedback (0-1) from past outcomes.
          Missing feedback defaults to 0.5.
        - **recency** (0.10): exponential decay favouring recent decisions.
        """
        now = time.time()

        # Group neighbours by (model, backend)
        route_groups: dict[tuple[str, str], list[tuple[float, EmbeddingMetadata]]] = {}
        for neighbor in neighbors:
            key = (neighbor.metadata.model_used, neighbor.metadata.backend_type)
            route_groups.setdefault(key, []).append((neighbor.score, neighbor.metadata))

        # Compute global normalisation bounds across all neighbours
        all_latencies = [n.metadata.latency_ms for n in neighbors if n.metadata.latency_ms > 0]
        all_costs = [n.metadata.cost for n in neighbors if n.metadata.cost > 0]
        max_latency = max(all_latencies) if all_latencies else 1.0
        max_cost = max(all_costs) if all_costs else 1.0
        total_neighbors = len(neighbors)

        scored: list[ScoredRoute] = []

        for (model, backend), entries in route_groups.items():
            scores_list = [s for s, _ in entries]
            metas = [m for _, m in entries]

            # --- model_fit: mean similarity * route frequency fraction ---
            mean_similarity = sum(scores_list) / len(scores_list)
            frequency = len(entries) / total_neighbors
            model_fit = mean_similarity * (0.5 + 0.5 * frequency)
            # Clamp to [0, 1]
            model_fit = max(0.0, min(1.0, model_fit))

            # --- latency: normalised inverse (lower latency = higher score) ---
            avg_latency = sum(m.latency_ms for m in metas) / len(metas) if metas else 0.0
            latency_score = 1.0 - avg_latency / max_latency if max_latency > 0 and avg_latency > 0 else 1.0

            # --- cost: normalised inverse (lower cost = higher score) ---
            avg_cost = sum(m.cost for m in metas) / len(metas) if metas else 0.0
            cost_score = 1.0 - avg_cost / max_cost if max_cost > 0 and avg_cost > 0 else 1.0

            # --- quality: average user feedback (default 0.5 if absent) ---
            feedbacks = [m.user_feedback for m in metas if m.user_feedback is not None]
            quality_score = sum(feedbacks) / len(feedbacks) if feedbacks else 0.5

            # --- recency: exponential decay (half-life = 1 day) ---
            half_life = 86400.0  # 24 hours in seconds
            recency_values: list[float] = []
            for m in metas:
                age = max(now - m.timestamp, 0.0)
                decay = 2.0 ** (-age / half_life)
                recency_values.append(decay)
            recency_score = sum(recency_values) / len(recency_values) if recency_values else 0.5

            # --- Composite weighted score ---
            factor_scores = {
                "model_fit": model_fit,
                "latency": latency_score,
                "cost": cost_score,
                "quality": quality_score,
                "recency": recency_score,
            }

            composite = sum(_FACTOR_WEIGHTS[factor] * value for factor, value in factor_scores.items())
            # Clamp to [0, 1]
            composite = max(0.0, min(1.0, composite))

            scored.append(
                ScoredRoute(
                    model=model,
                    backend=backend,
                    score=composite,
                    factor_scores=factor_scores,
                )
            )

        # Sort descending by composite score
        scored.sort(key=lambda r: r.score, reverse=True)
        return scored

    # ------------------------------------------------------------------
    # Outcome recording
    # ------------------------------------------------------------------

    async def record_outcome(
        self,
        request_id: str,
        embedding: np.ndarray,
        model_used: str,
        backend: str,
        latency_ms: float,
        cost: float,
        routing_tier: str = "knn",
        feedback: float | None = None,
        was_exploration: bool = False,
    ) -> None:
        """Record a routing outcome so future kNN lookups can learn from it.

        Parameters
        ----------
        request_id:
            Unique identifier for the request.
        embedding:
            Pre-computed embedding vector for the request text.
        model_used:
            The model that actually served the request.
        backend:
            The backend type that served the request.
        latency_ms:
            End-to-end latency in milliseconds.
        cost:
            Estimated cost in USD.
        routing_tier:
            Which tier made the final decision.
        feedback:
            Optional user quality rating (0.0-1.0).
        was_exploration:
            Whether this was a deliberate exploration request.
        """
        from llmhosts.embeddings.store import EmbeddingMetadata

        metadata = EmbeddingMetadata(
            request_id=request_id,
            model_used=model_used,
            backend_type=backend,
            routing_tier=routing_tier,
            user_feedback=feedback,
            latency_ms=latency_ms,
            cost=cost,
            was_exploration=was_exploration,
        )

        await self._store.add(embedding, metadata)

        # Periodically auto-save (every 100 additions)
        if self._store.size % 100 == 0:
            await self._store.save()
            logger.debug("Auto-saved FAISS index at %d embeddings", self._store.size)

    async def update_feedback(self, request_id: str, feedback: float) -> bool:
        """Update user feedback for a previously recorded outcome.

        Returns ``True`` if the record was found and updated.
        """
        return await self._store.update_feedback(request_id, feedback)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _format_factors(factors: dict[str, float]) -> str:
        """Format factor scores for human-readable reasoning strings."""
        parts = [f"{k}={v:.2f}" for k, v in sorted(factors.items())]
        return ", ".join(parts)

    @property
    def is_ready(self) -> bool:
        """``True`` when the pipeline and store are both initialised."""
        return self._pipeline.is_ready and self._store.is_ready
