"""
Embedding abstraction layer - wraps FastEmbed with sentence_transformers-compatible API.

Enables fast cold-start via ONNX while maintaining existing codebase compatibility.
FastEmbed loads in ~2-3 seconds vs sentence_transformers' ~32 seconds on Windows cold start.
"""

from typing import List, Union
import numpy as np


class Embedder:
    """
    Unified embedding interface compatible with existing ContextStore API.

    Uses FastEmbed (ONNX Runtime) for fast cold-start on Windows.
    API mimics sentence_transformers.SentenceTransformer for drop-in replacement.
    """

    # Model name mapping: short name -> FastEmbed full name
    # Note: all-mpnet-base-v2 is NOT supported by FastEmbed, use all-MiniLM-L6-v2 instead
    MODEL_MAP = {
        "all-MiniLM-L6-v2": "sentence-transformers/all-MiniLM-L6-v2",
        "all-MiniLM-L12-v2": "sentence-transformers/all-MiniLM-L12-v2",
        "paraphrase-multilingual-mpnet-base-v2": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
    }

    # Known dimensions for supported models
    DIM_MAP = {
        "sentence-transformers/all-MiniLM-L6-v2": 384,
        "sentence-transformers/all-MiniLM-L12-v2": 384,
        "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": 768,
    }

    def __init__(self, model_name: str, device: str = "cpu"):
        """
        Initialize embedding model.

        Args:
            model_name: Model name (sentence_transformers format, e.g., "all-MiniLM-L6-v2")
            device: "cuda" or "cpu" - Note: FastEmbed uses CPU by default (fast enough for embeddings)
        """
        from fastembed import TextEmbedding

        # Map model name to FastEmbed format
        fastembed_name = self.MODEL_MAP.get(model_name, f"sentence-transformers/{model_name}")

        # FastEmbed on CPU is fast enough for embedding models
        # CUDA would require onnxruntime-gpu which adds complexity
        # For small models like MiniLM, CPU is typically sufficient
        self._model = TextEmbedding(model_name=fastembed_name)
        self._model_name = fastembed_name
        self._device = "cpu"  # FastEmbed defaults to CPU
        self._dim = self.DIM_MAP.get(fastembed_name, 384)  # Default to 384 (MiniLM)

    def encode(
        self,
        texts: Union[str, List[str]],
        batch_size: int = 32,
        show_progress_bar: bool = False,
    ) -> np.ndarray:
        """
        Encode text(s) to embeddings.

        Matches sentence_transformers.SentenceTransformer.encode() signature.

        Args:
            texts: Single string or list of strings
            batch_size: Batch size for encoding (passed to FastEmbed)
            show_progress_bar: Ignored (FastEmbed doesn't support this)

        Returns:
            np.ndarray: Single embedding (1D) if input was string,
                       or array of embeddings (2D) if input was list
        """
        # Handle single string
        if isinstance(texts, str):
            embeddings = list(self._model.embed([texts], batch_size=batch_size))
            return np.array(embeddings[0])

        # Handle list of strings
        embeddings = list(self._model.embed(texts, batch_size=batch_size))
        return np.array(embeddings)

    def get_sentence_embedding_dimension(self) -> int:
        """Return embedding dimension for current model."""
        return self._dim
