"""Embedding generation for RAG pipelines."""

from __future__ import annotations

from abc import ABC, abstractmethod


class Embedder(ABC):
    """Base class for embedding providers."""

    @abstractmethod
    def embed(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings for a batch of texts."""
        ...

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


class VoyageEmbedder(Embedder):
    """Embedding provider using Voyage AI (via Anthropic partnership).

    Requires VOYAGE_API_KEY environment variable or explicit api_key.
    """

    DIMENSIONS = {
        "voyage-3": 1024,
        "voyage-3-lite": 512,
        "voyage-code-3": 1024,
    }

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "voyage-3",
        batch_size: int = 128,
    ):
        self.model = model
        self.batch_size = batch_size
        self.dimension = self.DIMENSIONS.get(model, 1024)
        self._api_key = api_key

    def embed(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings for a batch of texts.

        Requires voyageai package and valid API key at runtime.
        """
        import voyageai

        client = voyageai.Client(api_key=self._api_key)
        all_embeddings = []

        for i in range(0, len(texts), self.batch_size):
            batch = texts[i:i + self.batch_size]
            result = client.embed(batch, model=self.model, input_type="document")
            all_embeddings.extend(result.embeddings)

        return all_embeddings

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

        client = voyageai.Client(api_key=self._api_key)
        result = client.embed([text], model=self.model, input_type="query")
        return result.embeddings[0]
