"""
Embedding service for generating text embeddings.

Uses sentence-transformers with the all-MiniLM-L6-v2 model (384 dimensions).
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from reviewiq_pipeline.config import Config

logger = logging.getLogger(__name__)


class EmbeddingService:
    """
    Service for generating text embeddings using sentence-transformers.

    Uses the all-MiniLM-L6-v2 model by default, which produces 384-dimensional
    embeddings suitable for semantic similarity and clustering.
    """

    def __init__(self, config: Config):
        """
        Initialize the embedding service.

        Args:
            config: Pipeline configuration with embedding model settings
        """
        self.config = config
        self.model_name = config.embedding_model
        self.dimension = config.embedding_dimension
        self._model = None
        self._initialized = False

    def _ensure_initialized(self) -> None:
        """Lazy initialization of the sentence-transformers model."""
        if self._initialized:
            return

        try:
            from sentence_transformers import SentenceTransformer

            logger.info(f"Loading embedding model: {self.model_name}")
            self._model = SentenceTransformer(self.model_name)
            self._initialized = True
            logger.info(f"Embedding model loaded. Dimension: {self._model.get_sentence_embedding_dimension()}")
        except ImportError:
            raise ImportError(
                "sentence-transformers is required for embeddings. "
                "Install with: pip install sentence-transformers"
            )

    def embed(self, text: str) -> list[float]:
        """
        Generate embedding for a single text.

        Args:
            text: Text to embed

        Returns:
            List of floats representing the embedding vector
        """
        self._ensure_initialized()

        if not text or not text.strip():
            # Return zero vector for empty text
            return [0.0] * self.dimension

        embedding = self._model.encode(text, convert_to_numpy=True)
        return embedding.tolist()

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        """
        Generate embeddings for multiple texts.

        More efficient than calling embed() repeatedly.

        Args:
            texts: List of texts to embed

        Returns:
            List of embedding vectors
        """
        self._ensure_initialized()

        if not texts:
            return []

        # Handle empty strings
        non_empty_indices = [i for i, t in enumerate(texts) if t and t.strip()]
        non_empty_texts = [texts[i] for i in non_empty_indices]

        if not non_empty_texts:
            return [[0.0] * self.dimension for _ in texts]

        # Batch encode
        embeddings = self._model.encode(non_empty_texts, convert_to_numpy=True)

        # Build result with zero vectors for empty strings
        result = [[0.0] * self.dimension for _ in texts]
        for idx, emb in zip(non_empty_indices, embeddings):
            result[idx] = emb.tolist()

        return result

    def similarity(self, embedding1: list[float], embedding2: list[float]) -> float:
        """
        Calculate cosine similarity between two embeddings.

        Args:
            embedding1: First embedding vector
            embedding2: Second embedding vector

        Returns:
            Cosine similarity score between -1 and 1
        """
        vec1 = np.array(embedding1)
        vec2 = np.array(embedding2)

        # Handle zero vectors
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)

        if norm1 == 0 or norm2 == 0:
            return 0.0

        return float(np.dot(vec1, vec2) / (norm1 * norm2))

    def find_similar(
        self,
        query_embedding: list[float],
        candidate_embeddings: list[list[float]],
        top_k: int = 5,
        threshold: float = 0.0,
    ) -> list[tuple[int, float]]:
        """
        Find most similar embeddings to a query.

        Args:
            query_embedding: Query embedding vector
            candidate_embeddings: List of candidate embeddings
            top_k: Number of top results to return
            threshold: Minimum similarity threshold

        Returns:
            List of (index, similarity) tuples, sorted by similarity descending
        """
        if not candidate_embeddings:
            return []

        query = np.array(query_embedding)
        candidates = np.array(candidate_embeddings)

        # Compute all similarities at once
        query_norm = np.linalg.norm(query)
        if query_norm == 0:
            return []

        candidate_norms = np.linalg.norm(candidates, axis=1)

        # Avoid division by zero
        valid_mask = candidate_norms > 0
        similarities = np.zeros(len(candidates))
        similarities[valid_mask] = (
            np.dot(candidates[valid_mask], query)
            / (candidate_norms[valid_mask] * query_norm)
        )

        # Filter by threshold and get top k
        results = [
            (i, float(sim))
            for i, sim in enumerate(similarities)
            if sim >= threshold
        ]
        results.sort(key=lambda x: x[1], reverse=True)

        return results[:top_k]

    @property
    def model(self):
        """Get the underlying sentence-transformers model."""
        self._ensure_initialized()
        return self._model


def normalize_embedding(embedding: list[float]) -> list[float]:
    """
    Normalize an embedding to unit length.

    Args:
        embedding: Embedding vector

    Returns:
        Unit-normalized embedding
    """
    vec = np.array(embedding)
    norm = np.linalg.norm(vec)
    if norm == 0:
        return embedding
    return (vec / norm).tolist()


def average_embeddings(embeddings: list[list[float]]) -> list[float]:
    """
    Compute the average of multiple embeddings.

    Useful for creating centroid vectors for clustering.

    Args:
        embeddings: List of embedding vectors

    Returns:
        Averaged embedding vector
    """
    if not embeddings:
        raise ValueError("Cannot average empty embedding list")

    arr = np.array(embeddings)
    return arr.mean(axis=0).tolist()
