"""Embedding provider abstraction layer."""
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from numpy.typing import NDArray


class EmbeddingProvider(ABC):
    """Abstract base class for embedding providers."""

    @abstractmethod
    def encode(self, text: str) -> list[float]:
        """Encode text to embedding vector."""
        pass

    @abstractmethod
    def encode_batch(self, texts: list[str]) -> list[list[float]]:
        """Encode multiple texts to embedding vectors."""
        pass

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


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

    def __init__(self, model_name: str = "BAAI/bge-small-zh-v1.5"):
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError:
            raise ImportError(
                "sentence-transformers not installed. "
                "Install with: pip install sentence-transformers"
            )
        
        self._model = SentenceTransformer(model_name)
        self._model_name = model_name

    def encode(self, text: str) -> list[float]:
        embedding = self._model.encode(text, normalize_embeddings=True)
        return embedding.tolist()

    def encode_batch(self, texts: list[str]) -> list[list[float]]:
        embeddings = self._model.encode(texts, normalize_embeddings=True)
        return embeddings.tolist()

    @property
    def dimension(self) -> int:
        return self._model.get_sentence_embedding_dimension()


class OpenAIEmbedding(EmbeddingProvider):
    """OpenAI embedding API provider."""

    def __init__(self, api_key: str, model: str = "text-embedding-3-small"):
        try:
            from openai import OpenAI
        except ImportError:
            raise ImportError(
                "openai not installed. "
                "Install with: pip install openai"
            )
        
        self._client = OpenAI(api_key=api_key)
        self._model = model
        self._dimension = self._get_dimension(model)

    def _get_dimension(self, model: str) -> int:
        # text-embedding-3-small: 1536, text-embedding-3-large: 3072
        if "3-large" in model:
            return 3072
        return 1536

    def encode(self, text: str) -> list[float]:
        response = self._client.embeddings.create(
            model=self._model,
            input=text
        )
        return response.data[0].embedding

    def encode_batch(self, texts: list[str]) -> list[list[float]]:
        response = self._client.embeddings.create(
            model=self._model,
            input=texts
        )
        return [item.embedding for item in response.data]

    @property
    def dimension(self) -> int:
        return self._dimension


class DashscopeEmbedding(EmbeddingProvider):
    """Alibaba Dashscope embedding API provider."""

    def __init__(self, api_key: str, model: str = "text-embedding-v3"):
        try:
            import dashscope
        except ImportError:
            raise ImportError(
                "dashscope not installed. "
                "Install with: pip install dashscope"
            )
        
        dashscope.api_key = api_key
        self._model = model
        self._dimension = 1536  # v3 default

    def encode(self, text: str) -> list[float]:
        from dashscope import TextEmbedding
        
        response = TextEmbedding.call(
            model=self._model,
            input=text
        )
        return response.output['embeddings'][0]['embedding']

    def encode_batch(self, texts: list[str]) -> list[list[float]]:
        from dashscope import TextEmbedding
        
        response = TextEmbedding.call(
            model=self._model,
            input=texts
        )
        return [item['embedding'] for item in response.output['embeddings']]

    @property
    def dimension(self) -> int:
        return self._dimension


def create_embedding_provider(
    provider: str = "local",
    api_key: str | None = None,
    model: str = "BAAI/bge-small-zh-v1.5",
) -> EmbeddingProvider:
    """Factory function to create embedding provider."""
    
    provider = provider.lower()
    
    if provider == "local":
        return LocalEmbedding(model)
    elif provider == "openai":
        if not api_key:
            raise ValueError("api_key required for OpenAI embedding")
        return OpenAIEmbedding(api_key, model)
    elif provider == "dashscope":
        if not api_key:
            raise ValueError("api_key required for Dashscope embedding")
        return DashscopeEmbedding(api_key, model)
    else:
        raise ValueError(f"Unknown embedding provider: {provider}")
