import torch
from sentence_transformers import SentenceTransformer
import numpy as np
from typing import List, Union, Optional
import logging


class SentenceTransformerBasedEmbedder:
    """
    A class to generate embeddings from text using sentence transformers.
    Supports both CPU and GPU processing.
    """

    def __init__(
        self,
        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
        device: Optional[str] = None,
        use_gpu: bool = True,
        batch_size: int = 1,
        normalize_embeddings: bool = True,
    ):
        """
        Initialize the SentenceTransformerBasedEmbedder with a specific model.

        Args:
            model_name: The name of the sentence transformer model to use
            device: Manually specify the device to use ('cpu', 'cuda:0', etc.)
            use_gpu: Whether to use GPU if available
            batch_size: Batch size for processing multiple texts
            normalize_embeddings: Whether to L2-normalize the embeddings
        """
        self.logger = logging.getLogger(__name__)
        self.model_name = model_name
        self.batch_size = batch_size
        self.normalize_embeddings = normalize_embeddings

        # Determine the device to use
        if device is not None:
            self.device = device
        elif use_gpu and torch.cuda.is_available():
            self.device = "cuda"
            self.logger.info(f"Using GPU: {torch.cuda.get_device_name(0)}")
        else:
            if use_gpu and not torch.cuda.is_available():
                self.logger.warning(
                    "GPU requested but not available, falling back to CPU"
                )
            self.device = "cpu"
            self.logger.info("Using CPU for embeddings")

        # Load the model
        try:
            self.model = SentenceTransformer(model_name, device=self.device)
            embedding_dim = self.model.get_sentence_embedding_dimension()
            if embedding_dim is None:
                raise ValueError("Unable to retrieve sentence embedding dimension")
            self.embedding_dim: int = embedding_dim
            self.logger.info(
                f"Loaded model {model_name} with embedding dimension {self.embedding_dim}"
            )
        except Exception as e:
            self.logger.error(f"Failed to load model {model_name}: {str(e)}")
            raise

    def embed(
        self,
        texts: Union[str, List[str]],
        batch_size: Optional[int] = None,
        show_progress_bar: bool = False,
    ) -> np.ndarray:
        """
        Generate embeddings for the given text or list of texts.

        Args:
            texts: Text or list of texts to embed
            batch_size: Override the default batch size if specified
            show_progress_bar: Whether to show a progress bar for large batches

        Returns:
            Array of embeddings
        """
        if not texts:
            return np.array([])

        # Handle single text input
        if isinstance(texts, str):
            texts = [texts]

        # Use class batch size if not specified
        if batch_size is None:
            batch_size = self.batch_size

        try:
            # Generate embeddings
            embeddings = self.model.encode(
                texts,
                batch_size=batch_size,
                show_progress_bar=show_progress_bar,
                convert_to_numpy=True,
                normalize_embeddings=self.normalize_embeddings,
            )
            return embeddings
        except Exception as e:
            self.logger.error(f"Error generating embeddings: {str(e)}")
            raise

    def embed_documents(
        self, documents: List[dict], text_field: str = "text"
    ) -> List[dict]:
        """
        Generate embeddings for a list of document dictionaries.

        Args:
            documents: List of document dictionaries
            text_field: The field in each document containing the text to embed

        Returns:
            List of documents with 'embedding' field added
        """
        if not documents:
            return []

        # Extract texts from documents
        texts = [doc.get(text_field, "") for doc in documents]

        # Generate embeddings
        embeddings = self.embed(texts)

        # Add embeddings to documents
        for i, doc in enumerate(documents):
            doc["embedding"] = embeddings[i].tolist()

        return documents

    def get_embedding_dimension(self) -> int:
        """Get the dimension of the embeddings produced by the model"""
        return self.embedding_dim
