"""
Embedding Model Wrapper

Provides a unified interface for generating embeddings using
OpenAI or local models (BGE, sentence-transformers).
"""

import hashlib
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum

from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

logger = logging.getLogger(__name__)


class EmbeddingProvider(Enum):
    """Supported embedding providers."""

    OPENAI = "openai"
    BGE = "bge"
    SENTENCE_TRANSFORMERS = "sentence_transformers"


@dataclass
class EmbeddingConfig:
    """Configuration for embedding model."""

    provider: EmbeddingProvider = EmbeddingProvider.OPENAI
    model_name: str = "text-embedding-3-small"
    dimension: int = 1536

    # OpenAI specific
    api_key: str | None = None
    organization: str | None = None

    # Local model specific
    device: str = "cpu"  # or "cuda", "mps"
    normalize: bool = True

    # Batching
    batch_size: int = 100
    max_tokens: int = 8191  # OpenAI limit


class EmbeddingModel(ABC):
    """Abstract base class for embedding models."""

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

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

    @property
    @abstractmethod
    def dimension(self) -> int:
        """Return embedding dimension."""
        pass


class OpenAIEmbeddingError(Exception):
    """Base exception for OpenAI embedding errors."""

    pass


class OpenAIRateLimitError(OpenAIEmbeddingError):
    """Rate limit exceeded error."""

    pass


class OpenAIAPIError(OpenAIEmbeddingError):
    """General API error."""

    pass


class OpenAIEmbeddings(EmbeddingModel):
    """
    OpenAI embedding model wrapper with caching and retry logic.

    Uses text-embedding-3-small by default (1536 dimensions).

    Features:
        - Automatic retries with exponential backoff for rate limiting
        - In-memory caching to avoid redundant API calls
        - Efficient batching for multiple texts
        - Proper error handling for API errors
    """

    def __init__(self, config: EmbeddingConfig | None = None):
        """
        Initialize OpenAI embedding model.

        Args:
            config: Embedding configuration
        """
        self.config = config or EmbeddingConfig()
        self._client = None
        # Simple dict cache: hash(text) -> embedding
        self._cache: dict[str, list[float]] = {}

    def _get_client(self):
        """Get or create OpenAI async client."""
        if self._client is None:
            try:
                from openai import AsyncOpenAI
            except ImportError:
                raise ImportError("openai not installed. Run: pip install openai") from None

            self._client = AsyncOpenAI(
                api_key=self.config.api_key,
                organization=self.config.organization,
            )
        return self._client

    def _cache_key(self, text: str) -> str:
        """Generate a cache key for the given text."""
        # Include model name in key so different models don't share cache
        key_input = f"{self.config.model_name}:{text}"
        return hashlib.sha256(key_input.encode()).hexdigest()

    def _get_cached(self, text: str) -> list[float] | None:
        """Get embedding from cache if available."""
        key = self._cache_key(text)
        return self._cache.get(key)

    def _set_cached(self, text: str, embedding: list[float]) -> None:
        """Store embedding in cache."""
        key = self._cache_key(text)
        self._cache[key] = embedding

    def clear_cache(self) -> None:
        """Clear the embedding cache."""
        self._cache.clear()
        logger.debug("Embedding cache cleared")

    @property
    def cache_size(self) -> int:
        """Return the number of cached embeddings."""
        return len(self._cache)

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

        Args:
            text: The text to embed

        Returns:
            List of floats representing the embedding vector

        Raises:
            OpenAIEmbeddingError: If the API call fails
        """
        # Check cache first
        cached = self._get_cached(text)
        if cached is not None:
            logger.debug("Cache hit for embedding")
            return cached

        embeddings = await self.embed_batch([text])
        return embeddings[0]

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

        Args:
            texts: List of texts to embed

        Returns:
            List of embedding vectors (same order as input texts)

        Raises:
            OpenAIEmbeddingError: If the API call fails
        """
        if not texts:
            return []

        # Separate cached and uncached texts
        results: dict[int, list[float]] = {}
        uncached_indices: list[int] = []
        uncached_texts: list[str] = []

        for i, text in enumerate(texts):
            cached = self._get_cached(text)
            if cached is not None:
                results[i] = cached
            else:
                uncached_indices.append(i)
                uncached_texts.append(text)

        if uncached_texts:
            logger.debug(
                f"Embedding {len(uncached_texts)} texts "
                f"({len(results)} cached)"
            )
            # Fetch uncached embeddings from API
            new_embeddings = await self._fetch_embeddings(uncached_texts)

            # Store in cache and results
            for idx, text, embedding in zip(
                uncached_indices, uncached_texts, new_embeddings, strict=True
            ):
                self._set_cached(text, embedding)
                results[idx] = embedding
        else:
            logger.debug(f"All {len(texts)} embeddings served from cache")

        # Return in original order
        return [results[i] for i in range(len(texts))]

    async def _fetch_embeddings(self, texts: list[str]) -> list[list[float]]:
        """
        Fetch embeddings from OpenAI API with batching and retries.

        Args:
            texts: List of texts to embed (should not be empty)

        Returns:
            List of embedding vectors
        """
        client = self._get_client()
        all_embeddings: list[list[float]] = []

        # Process in batches
        for i in range(0, len(texts), self.config.batch_size):
            batch = texts[i : i + self.config.batch_size]

            # Truncate texts that are too long (rough estimate: 4 chars per token)
            truncated_batch = [
                text[: self.config.max_tokens * 4] for text in batch
            ]

            batch_embeddings = await self._call_api_with_retry(
                client, truncated_batch
            )
            all_embeddings.extend(batch_embeddings)

        return all_embeddings

    @retry(
        retry=retry_if_exception_type(OpenAIRateLimitError),
        wait=wait_exponential(multiplier=1, min=1, max=60),
        stop=stop_after_attempt(5),
        reraise=True,
    )
    async def _call_api_with_retry(
        self, client, texts: list[str]
    ) -> list[list[float]]:
        """
        Call OpenAI API with automatic retry on rate limits.

        Args:
            client: OpenAI async client
            texts: List of texts to embed

        Returns:
            List of embedding vectors

        Raises:
            OpenAIRateLimitError: If rate limit is hit (will be retried)
            OpenAIAPIError: For other API errors
        """
        try:
            from openai import APIError, RateLimitError
        except ImportError:
            raise ImportError("openai not installed. Run: pip install openai") from None

        try:
            response = await client.embeddings.create(
                input=texts,
                model=self.config.model_name,
            )

            # Extract embeddings in the correct order
            # API returns data sorted by index
            sorted_data = sorted(response.data, key=lambda x: x.index)
            return [item.embedding for item in sorted_data]

        except RateLimitError as e:
            logger.warning(f"Rate limit hit, will retry: {e}")
            raise OpenAIRateLimitError(str(e)) from e

        except APIError as e:
            logger.error(f"OpenAI API error: {e}")
            raise OpenAIAPIError(f"API error: {e}") from e

        except Exception as e:
            logger.error(f"Unexpected error calling OpenAI: {e}")
            raise OpenAIEmbeddingError(f"Unexpected error: {e}") from e

    @property
    def dimension(self) -> int:
        """Return embedding dimension."""
        return self.config.dimension


# Alias for backward compatibility
OpenAIEmbedding = OpenAIEmbeddings


class LocalEmbedding(EmbeddingModel):
    """
    Local embedding model using sentence-transformers.

    Supports models like BGE-large-en-v1.5 (1024 dimensions).
    """

    def __init__(self, config: EmbeddingConfig | None = None):
        """
        Initialize local embedding model.

        Args:
            config: Embedding configuration
        """
        self.config = config or EmbeddingConfig(
            provider=EmbeddingProvider.BGE,
            model_name="BAAI/bge-large-en-v1.5",
            dimension=1024,
        )
        self._model = None

    def _get_model(self):
        """Get or create sentence-transformer model."""
        if self._model is None:
            try:
                from sentence_transformers import SentenceTransformer
            except ImportError:
                raise ImportError(
                    "sentence-transformers not installed. "
                    "Run: pip install sentence-transformers"
                ) from None

            self._model = SentenceTransformer(
                self.config.model_name,
                device=self.config.device,
            )
        return self._model

    async def embed(self, text: str) -> list[float]:
        """Generate embedding for a single text."""
        embeddings = await self.embed_batch([text])
        return embeddings[0]

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

        try:
            embeddings = model.encode(
                texts,
                batch_size=self.config.batch_size,
                normalize_embeddings=self.config.normalize,
                show_progress_bar=False,
            )
            return embeddings.tolist()

        except Exception as e:
            logger.error(f"Local embedding error: {e}")
            raise

    @property
    def dimension(self) -> int:
        """Return embedding dimension."""
        return self.config.dimension


class SparseEncoder:
    """
    Sparse encoder for BM25-style vectors.

    Uses fastembed for efficient sparse encoding.
    """

    def __init__(self, model_name: str = "prithivida/Splade_PP_en_v1"):
        """
        Initialize sparse encoder.

        Args:
            model_name: Name of sparse encoder model
        """
        self.model_name = model_name
        self._model = None

    def _get_model(self):
        """Get or create sparse encoder model."""
        if self._model is None:
            try:
                from fastembed import SparseTextEmbedding
            except ImportError:
                raise ImportError(
                    "fastembed not installed. Run: pip install fastembed"
                ) from None

            self._model = SparseTextEmbedding(model_name=self.model_name)
        return self._model

    async def encode(self, text: str) -> dict:
        """
        Encode text to sparse vector.

        Returns:
            Dict with 'indices' and 'values' keys
        """
        embeddings = await self.encode_batch([text])
        return embeddings[0]

    async def encode_batch(self, texts: list[str]) -> list[dict]:
        """
        Encode multiple texts to sparse vectors.

        Returns:
            List of dicts with 'indices' and 'values' keys
        """
        model = self._get_model()

        try:
            results = []
            embeddings = list(model.embed(texts))

            for emb in embeddings:
                results.append({
                    "indices": emb.indices.tolist(),
                    "values": emb.values.tolist(),
                })

            return results

        except Exception as e:
            logger.error(f"Sparse encoding error: {e}")
            raise


def create_embedding_model(config: EmbeddingConfig | None = None) -> EmbeddingModel:
    """
    Factory function to create embedding model.

    Args:
        config: Embedding configuration

    Returns:
        EmbeddingModel instance
    """
    config = config or EmbeddingConfig()

    if config.provider == EmbeddingProvider.OPENAI:
        return OpenAIEmbedding(config)
    elif config.provider in [EmbeddingProvider.BGE, EmbeddingProvider.SENTENCE_TRANSFORMERS]:
        return LocalEmbedding(config)
    else:
        raise ValueError(f"Unknown embedding provider: {config.provider}")
