# Copyright (C) 2026 GoodData Corporation

"""Embedding model wrapper using fastembed (ONNX Runtime)."""

from __future__ import annotations

import logging
import os
from pathlib import Path

import numpy as np
from fastembed import TextEmbedding
from numpy.typing import NDArray

logger = logging.getLogger(__name__)

# Default model - small, fast, good for semantic search
# fastembed requires the full HuggingFace repo ID
DEFAULT_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"

# Known embedding dimension for all-MiniLM-L6-v2
_KNOWN_DIMENSIONS: dict[str, int] = {
    "sentence-transformers/all-MiniLM-L6-v2": 384,
}


def _default_fastembed_cache_dir() -> Path:
    raw = (
        os.environ.get("AIDA_FASTEMBED_CACHE_DIR")
        or os.environ.get("FASTEMBED_CACHE_DIR")
        or os.environ.get("FASTEMBED_CACHE_PATH")
    )
    if raw:
        return Path(raw).expanduser()
    # Prefer fastembed's documented default location.
    return Path.home() / ".cache" / "fastembed"


def is_model_cached(model_name: str = DEFAULT_MODEL_NAME) -> bool:
    """
    Best-effort check whether a fastembed model is present in local cache.

    fastembed stores models under ~/.cache/fastembed/ by default,
    in a directory derived from the model name.
    """
    cache_dir = _default_fastembed_cache_dir()
    if not cache_dir.is_dir():
        return False

    # fastembed cache directory naming varies across versions:
    #   - older: "sentence-transformers--all-MiniLM-L6-v2"  (HF repo ID, / -> --)
    #   - newer: "models--qdrant--all-MiniLM-L6-v2-onnx"   (qdrant ONNX variant)
    # Match on the base model name (last path segment) to handle both conventions.
    base_name = model_name.rsplit("/", 1)[-1]  # e.g. "all-MiniLM-L6-v2"
    try:
        return any(base_name in d.name for d in cache_dir.iterdir() if d.is_dir())
    except OSError:
        return False


class EmbeddingModel:
    """Wrapper for fastembed embedding model."""

    def __init__(self, model_name: str = DEFAULT_MODEL_NAME, cache_dir: Path | None = None):
        """
        Initialize the embedding model.

        Args:
            model_name: Name of the fastembed model (full HuggingFace repo ID)
            cache_dir: Directory to cache the model (uses default if None)
        """
        self.model_name = model_name
        self._model: TextEmbedding | None = None
        self._cache_dir = cache_dir
        self._dimension: int | None = _KNOWN_DIMENSIONS.get(model_name)

    def _load_model(self) -> TextEmbedding:
        """Lazy-load the model on first use."""
        if self._model is None:
            logger.info(f"Loading embedding model: {self.model_name}")
            kwargs: dict[str, object] = {"model_name": self.model_name}
            cache_dir = self._cache_dir or _default_fastembed_cache_dir()
            try:
                cache_dir.mkdir(parents=True, exist_ok=True)
                kwargs["cache_dir"] = str(cache_dir)
            except OSError:
                logger.warning("Unable to create fastembed cache dir; falling back to library default.")
            self._model = TextEmbedding(**kwargs)  # type: ignore[arg-type]

            # Detect dimension if not known
            if self._dimension is None:
                probe = list(self._model.embed(["dimension probe"]))
                self._dimension = probe[0].shape[0]

            logger.info(f"Loaded embedding model with dimension: {self._dimension}")

        return self._model

    @property
    def dimension(self) -> int:
        """Get the embedding dimension."""
        if self._dimension is None:
            self._load_model()
        assert self._dimension is not None
        return self._dimension

    def embed(self, texts: list[str]) -> NDArray[np.float32]:
        """
        Generate embeddings for a list of texts.

        Args:
            texts: List of text strings to embed

        Returns:
            Numpy array of shape (len(texts), dimension)
        """
        if not texts:
            return np.array([], dtype=np.float32).reshape(0, self.dimension)

        model = self._load_model()
        # fastembed returns a generator of numpy arrays (already L2-normalized)
        embeddings_list = list(model.embed(texts))
        embeddings: NDArray[np.float32] = np.array(embeddings_list, dtype=np.float32)
        return embeddings

    def embed_query(self, query: str) -> NDArray[np.float32]:
        """
        Generate embedding for a single query.

        Args:
            query: Query text

        Returns:
            Numpy array of shape (dimension,)
        """
        embeddings = self.embed([query])
        return embeddings[0]


# Global singleton for efficiency
_global_model: EmbeddingModel | None = None


def get_embedding_model(model_name: str = DEFAULT_MODEL_NAME) -> EmbeddingModel:
    """Get the global embedding model instance (creates if needed)."""
    global _global_model
    if _global_model is None or _global_model.model_name != model_name:
        _global_model = EmbeddingModel(model_name)
    return _global_model


def init_embedding_model(model_name: str = DEFAULT_MODEL_NAME) -> None:
    """Pre-initialize the embedding model (call at startup)."""
    model = get_embedding_model(model_name)
    # Force model loading
    _ = model.dimension
    logger.info("Embedding model initialized")
