"""Local Qwen3-Embedding-8B implementation for GNN router inference.

This module provides a local GPU-based embedding provider using the Qwen3-Embedding-8B
model with instruction-conditioned embeddings optimized for LLM routing.

The model uses a router-specific instruction that guides embeddings to cluster by:
- Task type (math, coding, reasoning, factual, creative)
- Difficulty level (easy, medium, hard)
- Required LLM capabilities (tool use, long context, domain expertise)

This differs from default retrieval embeddings which optimize for semantic similarity.

Example:
    embedder = LocalQwen3Embedding.load(device="cuda", dimension=256)
    embeddings = embedder.embed(["What is 2+2?", "Prove the Riemann hypothesis"])
"""

from __future__ import annotations

import time
from typing import TYPE_CHECKING

import numpy as np

from embergraphrouter.embeddings.base import EmbeddingProvider, EmbeddingResult

if TYPE_CHECKING:
    from sentence_transformers import SentenceTransformer

# Router-specific instruction for query embeddings.
# This instruction guides the model to encode routing-relevant features rather than
# semantic similarity for retrieval.
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 identifier on HuggingFace.
DEFAULT_MODEL_NAME = "Qwen/Qwen3-Embedding-8B"

# Default target dimension for Matryoshka truncation.
# The 256-dim embeddings match training distribution and balance quality vs speed.
DEFAULT_DIMENSION = 256


class LocalQwen3Embedding:
    """Local Qwen3-Embedding-8B provider for routing embeddings.

    This implementation loads the Qwen3-Embedding-8B model locally on GPU and generates
    instruction-conditioned embeddings suitable for GNN routing. It supports Matryoshka
    dimension truncation for efficiency while preserving embedding quality.

    The class follows the EmbeddingProvider protocol, allowing it to be swapped with
    remote implementations without changing routing logic.

    Attributes:
        model_name: HuggingFace model identifier.
        dimension: Output embedding dimension (Matryoshka truncation).
        device: Compute device (e.g., "cuda", "cuda:0", "cpu").
        instruction: Instruction prefix for query embeddings.

    Example:
        # Load with defaults (256-dim on CUDA)
        embedder = LocalQwen3Embedding.load()
        embeddings = embedder.embed(["What is 2+2?"])

        # Load with custom dimension
        embedder = LocalQwen3Embedding.load(dimension=512, device="cuda:1")
    """

    def __init__(
        self,
        model: "SentenceTransformer",
        model_name: str,
        dimension: int,
        device: str,
        instruction: str,
    ) -> None:
        """Initialize with a loaded model.

        Use LocalQwen3Embedding.load() instead of calling this directly.

        Args:
            model: Loaded SentenceTransformer model.
            model_name: HuggingFace model identifier.
            dimension: Target embedding dimension.
            device: Compute device.
            instruction: Instruction prefix for embeddings.
        """
        self._model = model
        self._model_name = model_name
        self._dimension = dimension
        self._device = device
        self._instruction = instruction
        self._native_dimension = model.get_sentence_embedding_dimension()

    @classmethod
    def load(
        cls,
        model_name: str = DEFAULT_MODEL_NAME,
        dimension: int = DEFAULT_DIMENSION,
        device: str = "cuda",
        instruction: str = ROUTER_INSTRUCTION,
    ) -> "LocalQwen3Embedding":
        """Load the Qwen3-Embedding-8B model.

        Args:
            model_name: HuggingFace model identifier.
            dimension: Target dimension for Matryoshka truncation. The model natively
                produces 4096-dim embeddings; we truncate to dimension and re-normalize.
            device: Compute device (e.g., "cuda", "cuda:0", "cpu").
            instruction: Instruction prefix for query embeddings.

        Returns:
            Initialized LocalQwen3Embedding instance.

        Raises:
            RuntimeError: If model loading fails.
        """
        try:
            import torch
            from sentence_transformers import SentenceTransformer
        except ImportError as e:
            raise RuntimeError(
                "sentence-transformers and torch are required for LocalQwen3Embedding. "
                "Install with: pip install sentence-transformers torch"
            ) from e

        # Load directly to GPU in bfloat16 to avoid CPU memory spike
        # 8B parameters * 2 bytes = ~16GB GPU memory
        # Using device_map="cuda" loads directly to GPU, skipping CPU staging
        model = SentenceTransformer(
            model_name,
            device=device,
            trust_remote_code=True,
            model_kwargs={
                "torch_dtype": torch.bfloat16,
                "device_map": device,  # Load directly to GPU
                "low_cpu_mem_usage": True,  # Don't duplicate in CPU RAM
            },
        )
        return cls(
            model=model,
            model_name=model_name,
            dimension=dimension,
            device=device,
            instruction=instruction,
        )

    @property
    def model_name(self) -> str:
        """Return the model identifier."""
        return self._model_name

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

    @property
    def device(self) -> str:
        """Return the compute device."""
        return self._device

    @property
    def instruction(self) -> str:
        """Return the instruction prefix."""
        return self._instruction

    @property
    def native_dimension(self) -> int:
        """Return the model's native embedding dimension before truncation."""
        return self._native_dimension

    def embed(self, texts: list[str], batch_size: int = 32) -> np.ndarray:
        """Embed a batch of texts into dense vectors.

        Each text is prefixed with the router instruction in Qwen3 format:
            "Instruct: {instruction}\\nQuery: {text}"

        The embeddings are L2 normalized, then truncated to the target dimension,
        then re-normalized (Matryoshka truncation).

        Args:
            texts: List of text strings to embed.
            batch_size: Batch size for inference.

        Returns:
            Array of shape (len(texts), self.dimension) containing normalized embeddings.

        Raises:
            RuntimeError: If the embedding operation fails.
        """
        if not texts:
            return np.zeros((0, self._dimension), dtype=np.float32)

        # Format texts with instruction prefix (Qwen3 format)
        formatted_texts = [
            f"Instruct: {self._instruction}\nQuery: {text}" for text in texts
        ]

        # Encode in batches
        embeddings_list = []
        for i in range(0, len(formatted_texts), batch_size):
            batch = formatted_texts[i : i + batch_size]
            batch_embeddings = self._model.encode(
                batch,
                convert_to_numpy=True,
                normalize_embeddings=True,
            )
            embeddings_list.append(batch_embeddings)

        embeddings = np.vstack(embeddings_list)

        # Apply Matryoshka truncation if needed
        if self._dimension < embeddings.shape[1]:
            embeddings = embeddings[:, : self._dimension]
            # Re-normalize after truncation
            norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
            embeddings = embeddings / np.maximum(norms, 1e-12)

        return embeddings.astype(np.float32)

    def embed_with_result(
        self, texts: list[str], batch_size: int = 32
    ) -> EmbeddingResult:
        """Embed texts and return detailed result with metadata.

        Args:
            texts: List of text strings to embed.
            batch_size: Batch size for inference.

        Returns:
            EmbeddingResult with embeddings, model info, and timing.
        """
        start_time = time.perf_counter()
        embeddings = self.embed(texts, batch_size=batch_size)
        elapsed_ms = (time.perf_counter() - start_time) * 1000

        return EmbeddingResult(
            embeddings=embeddings,
            model_name=self._model_name,
            dimension=self._dimension,
            latency_ms=elapsed_ms,
        )

    def embed_documents(
        self, texts: list[str], batch_size: int = 32
    ) -> np.ndarray:
        """Embed documents without instruction prefix.

        For symmetric document embeddings (e.g., LLM descriptions), the instruction
        prefix is omitted as per Qwen3 documentation.

        Args:
            texts: List of document texts to embed.
            batch_size: Batch size for inference.

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

        embeddings_list = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            batch_embeddings = self._model.encode(
                batch,
                convert_to_numpy=True,
                normalize_embeddings=True,
            )
            embeddings_list.append(batch_embeddings)

        embeddings = np.vstack(embeddings_list)

        if self._dimension < embeddings.shape[1]:
            embeddings = embeddings[:, : self._dimension]
            norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
            embeddings = embeddings / np.maximum(norms, 1e-12)

        return embeddings.astype(np.float32)


# Protocol compliance verification
def _verify_protocol_compliance() -> None:
    """Verify LocalQwen3Embedding implements EmbeddingProvider protocol."""
    # This is a compile-time check; runtime verification happens via isinstance
    instance: EmbeddingProvider = LocalQwen3Embedding.__new__(LocalQwen3Embedding)  # noqa: F841
