"""LRU caching and batch embedding for EB1-GNN routing.

This module provides LRU-cached, batch-aware embedding via Qwen3-Embedding-8B.

Throughput model (B200 GPU, bfloat16):
- Single query (uncached): ~133ms
- Single query (cached):   ~0ms
- Batch of N queries (uncached): ~133ms (GPU parallelism)
- embed_batch() gives ~Nx throughput vs sequential embed() calls

The batch path is the key throughput optimization: collect cache misses,
encode them in one model.encode() call, insert results into cache.

Note: FP16 precision was tested and found to be 29% SLOWER than bfloat16 on B200.
"""

from __future__ import annotations

import hashlib
import logging
import threading
from collections import OrderedDict
from typing import TYPE_CHECKING, Any

import numpy as np
import torch

if TYPE_CHECKING:
    from sentence_transformers import SentenceTransformer

logger = logging.getLogger(__name__)

ROUTER_INSTRUCTION = (
    "Classify this query by task type (reasoning, mathematics, coding, "
    "factual recall, creative writing), cognitive difficulty level, and "
    "the language model capabilities required to answer it accurately."
)

DEFAULT_MODEL_NAME = "Qwen/Qwen3-Embedding-8B"
DEFAULT_DIMENSION = 256


class OptimizedQwen3Embedding:
    """Qwen3 embedding provider with LRU caching and true batch support.

    The LRU cache uses an OrderedDict so embed_batch() can probe for cache
    hits without triggering single-query GPU compute on misses. This enables
    collecting all misses and encoding them in one batched GPU call.

    Measured performance (B200 GPU, Jan 2026):
    - Uncached single: ~133ms
    - Cached single:   ~0ms
    - Uncached batch of 16: ~133ms (same wall time, 16x throughput)

    Attributes:
        model_name: HuggingFace model identifier.
        dimension: Output embedding dimension.
        device: Compute device (cuda/cpu).
        cache_size: LRU cache size (0 disables caching).
    """

    def __init__(
        self,
        model: "SentenceTransformer",
        model_name: str,
        dimension: int,
        device: str,
        instruction: str,
        cache_size: int = 1000,
    ) -> None:
        """Initialize embedding provider with caching.

        Args:
            model: Loaded SentenceTransformer model.
            model_name: Model identifier.
            dimension: Target embedding dimension.
            device: Compute device.
            instruction: Instruction prefix.
            cache_size: LRU cache size (0 to disable).
        """
        self._model = model
        self._model_name = model_name
        self._dimension = dimension
        self._device = device
        self._instruction = instruction
        self._cache_size = cache_size
        self._native_dimension = model.get_sentence_embedding_dimension()

        self._cache: OrderedDict[str, np.ndarray] = OrderedDict()
        self._cache_lock = threading.Lock()
        self._hits = 0
        self._misses = 0

    @classmethod
    def load(
        cls,
        model_name: str = DEFAULT_MODEL_NAME,
        dimension: int = DEFAULT_DIMENSION,
        device: str = "cuda",
        instruction: str = ROUTER_INSTRUCTION,
        cache_size: int = 1000,
    ) -> "OptimizedQwen3Embedding":
        """Load Qwen3 embedding model with caching.

        Args:
            model_name: HuggingFace model identifier.
            dimension: Target dimension for Matryoshka truncation.
            device: Compute device (cuda/cpu).
            instruction: Instruction prefix.
            cache_size: LRU cache size, 0 disables (default: 1000).

        Returns:
            Initialized OptimizedQwen3Embedding instance.
        """
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError as e:
            raise RuntimeError(
                "sentence-transformers required. Install with: "
                "pip install sentence-transformers"
            ) from e

        model = SentenceTransformer(
            model_name,
            device=device,
            trust_remote_code=True,
            model_kwargs={
                "torch_dtype": torch.bfloat16,
                "device_map": device,
                "low_cpu_mem_usage": True,
            },
        )

        return cls(
            model=model,
            model_name=model_name,
            dimension=dimension,
            device=device,
            instruction=instruction,
            cache_size=cache_size,
        )

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

    @property
    def cache_info(self) -> dict[str, Any]:
        """Return cache statistics."""
        if self._cache_size == 0:
            return {"enabled": False}
        total = self._hits + self._misses
        return {
            "enabled": True,
            "size": self._cache_size,
            "current": len(self._cache),
            "hits": self._hits,
            "misses": self._misses,
            "hit_rate": self._hits / total if total > 0 else 0.0,
        }

    def _hash_query(self, text: str) -> str:
        """Hash query for cache key (instruction + text)."""
        full_text = f"{self._instruction}\n{text}"
        return hashlib.sha256(full_text.encode()).hexdigest()[:16]

    def _cache_get(self, key: str) -> np.ndarray | None:
        """Look up key in LRU cache, promoting to most-recent on hit."""
        if self._cache_size == 0:
            return None
        with self._cache_lock:
            if key in self._cache:
                self._cache.move_to_end(key)
                self._hits += 1
                return self._cache[key]
            self._misses += 1
            return None

    def _cache_put(self, key: str, value: np.ndarray) -> None:
        """Insert into LRU cache, evicting oldest if at capacity."""
        if self._cache_size == 0:
            return
        with self._cache_lock:
            self._cache[key] = value
            self._cache.move_to_end(key)
            if len(self._cache) > self._cache_size:
                self._cache.popitem(last=False)

    def _truncate_and_normalize(self, embedding: np.ndarray) -> np.ndarray:
        """Truncate to target dimension and renormalize (Matryoshka).

        Args:
            embedding: Array of shape (dimension,) or (N, dimension).

        Returns:
            Same shape, truncated and L2-normalized.
        """
        if self._dimension >= self._native_dimension:
            return embedding.astype(np.float32)

        is_batch = embedding.ndim == 2
        if is_batch:
            embedding = embedding[:, : self._dimension]
            norms = np.linalg.norm(embedding, axis=1, keepdims=True)
            norms = np.maximum(norms, 1e-12)
            embedding = embedding / norms
        else:
            embedding = embedding[: self._dimension]
            norm = np.linalg.norm(embedding)
            if norm > 0:
                embedding = embedding / norm
        return embedding.astype(np.float32)

    def _encode_single(self, text: str) -> np.ndarray:
        """Encode a single text on GPU (no caching)."""
        formatted = f"Instruct: {self._instruction}\nQuery: {text}"
        embedding = self._model.encode(
            formatted,
            convert_to_numpy=True,
            normalize_embeddings=True,
            batch_size=1,
        )
        return self._truncate_and_normalize(embedding)

    def _encode_batch(
        self, texts: list[str], batch_size: int = 32
    ) -> np.ndarray:
        """Encode multiple texts in one GPU call (no caching).

        Args:
            texts: Raw query texts (without instruction prefix).
            batch_size: Batch size for model.encode().

        Returns:
            Array of shape (len(texts), dimension).
        """
        formatted = [
            f"Instruct: {self._instruction}\nQuery: {text}"
            for text in texts
        ]
        raw = self._model.encode(
            formatted,
            convert_to_numpy=True,
            normalize_embeddings=True,
            batch_size=batch_size,
        )
        return self._truncate_and_normalize(raw)

    def embed(self, query: str) -> np.ndarray:
        """Embed a single query with automatic caching.

        This method satisfies the EmbeddingProvider protocol.

        Args:
            query: Text string to embed.

        Returns:
            Array of shape (dimension,) with normalized embedding.
        """
        key = self._hash_query(query)
        cached = self._cache_get(key)
        if cached is not None:
            return cached

        embedding = self._encode_single(query)
        self._cache_put(key, embedding)
        return embedding

    def embed_batch(self, texts: list[str], batch_size: int = 32) -> np.ndarray:
        """Embed multiple texts with cache-aware GPU batching.

        Cache hits are returned immediately. Cache misses are collected and
        batch-encoded in a single GPU call, then inserted into the cache.
        One model.encode() call for N queries takes roughly the same wall
        time as encoding 1 query -- this is the throughput multiplier.

        Args:
            texts: List of text strings to embed.
            batch_size: Batch size passed to model.encode() for misses.

        Returns:
            Array of shape (len(texts), dimension) with normalized embeddings.
        """
        if not texts:
            return np.zeros((0, self._dimension), dtype=np.float32)

        results: list[np.ndarray | None] = [None] * len(texts)
        miss_indices: list[int] = []
        miss_keys: list[str] = []
        miss_texts: list[str] = []

        # Phase 1: probe cache, collect misses
        for i, text in enumerate(texts):
            key = self._hash_query(text)
            cached = self._cache_get(key)
            if cached is not None:
                results[i] = cached
            else:
                miss_indices.append(i)
                miss_keys.append(key)
                miss_texts.append(text)

        # Phase 2: batch-encode all misses in one GPU call
        if miss_texts:
            batch_embeddings = self._encode_batch(miss_texts, batch_size)
            for j, idx in enumerate(miss_indices):
                embedding = batch_embeddings[j]
                results[idx] = embedding
                self._cache_put(miss_keys[j], embedding)

        return np.stack(results)  # type: ignore[arg-type]

    def clear_cache(self) -> None:
        """Clear the embedding cache and reset statistics."""
        with self._cache_lock:
            self._cache.clear()
            self._hits = 0
            self._misses = 0
