"""Embedding model wrapper using sentence-transformers."""

from __future__ import annotations

import contextlib
import logging
import os
from collections.abc import Iterator
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
from numpy.typing import NDArray
from sentence_transformers import SentenceTransformer

if TYPE_CHECKING:
    pass  # All types available at runtime

logger = logging.getLogger(__name__)

# Default model - small, fast, good for semantic search
# Same family as gen-ai service uses
DEFAULT_MODEL_NAME = "all-MiniLM-L6-v2"


def _normalize_hf_repo_id(model_name: str) -> str:
    """
    Normalize a sentence-transformers model name to a HuggingFace repo id.

    SentenceTransformer accepts both:
    - "all-MiniLM-L6-v2" (implicit sentence-transformers org)
    - "sentence-transformers/all-MiniLM-L6-v2" (explicit repo id)
    """
    if "/" in model_name:
        return model_name
    return f"sentence-transformers/{model_name}"


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

    This avoids triggering any network calls.
    """
    try:
        from huggingface_hub.constants import HF_HUB_CACHE
    except Exception:
        # If huggingface_hub isn't importable, we can't reliably detect cache state.
        return False

    repo_id = _normalize_hf_repo_id(model_name)
    org, repo = repo_id.split("/", 1)
    model_dir = Path(HF_HUB_CACHE) / f"models--{org}--{repo}"
    refs_dir = model_dir / "refs"
    if not refs_dir.is_dir():
        return False
    try:
        return any(refs_dir.iterdir())
    except OSError:
        return False


@contextlib.contextmanager
def _quiet_hf_download_and_load_noise() -> Iterator[None]:
    """
    Reduce third-party log/progress spam during model download/load.

    Cursor sometimes labels stderr/progress as "[error]" in the MCP log. We prefer to keep
    the first-run experience readable rather than dumping ML loader progress bars.
    """
    old_env = os.environ.get("HF_HUB_DISABLE_PROGRESS_BARS")
    os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
    try:
        # HuggingFace hub warnings (e.g. unauthenticated requests) are logged at WARNING by default.
        try:
            from huggingface_hub import logging as hf_logging

            hf_logging.set_verbosity_error()
        except Exception:
            pass

        # Transformers model loading uses tqdm progress bars unless disabled.
        try:
            from transformers.utils import logging as tf_logging

            tf_logging.set_verbosity_error()
            tf_logging.disable_progress_bar()
        except Exception:
            pass

        yield None
    finally:
        if old_env is None:
            os.environ.pop("HF_HUB_DISABLE_PROGRESS_BARS", None)
        else:
            os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = old_env


class EmbeddingModel:
    """Wrapper for sentence-transformers 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 sentence-transformers model
            cache_dir: Directory to cache the model (uses default if None)
        """
        self.model_name = model_name
        self._model: SentenceTransformer | None = None
        self._cache_dir = cache_dir
        self._dimension: int | None = None

    def _load_model(self) -> SentenceTransformer:
        """Lazy-load the model on first use."""
        if self._model is None:
            logger.info(f"Loading embedding model: {self.model_name}")
            with _quiet_hf_download_and_load_noise():
                self._model = SentenceTransformer(
                    self.model_name,
                    cache_folder=str(self._cache_dir) if self._cache_dir else None,
                    device="cpu",
                )
            self._dimension = self._model.get_sentence_embedding_dimension()
            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()
        # sentence-transformers has **kwargs: Unknown in encode() signature
        result = model.encode(
            texts,
            convert_to_numpy=True,
            show_progress_bar=False,
            normalize_embeddings=True,  # For cosine similarity
        )
        embeddings: NDArray[np.float32] = np.asarray(result, 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")
