from __future__ import annotations

from typing import List, Sequence

import anyio

from idxr.embeddings import SentenceTransformerBasedEmbedder

from openai import OpenAI
from openai.types import CreateEmbeddingResponse

from ..interfaces import EmbeddingStrategy


class SentenceTransformerEmbeddingStrategy(EmbeddingStrategy):
    """
    Local embedding backend powered by sentence-transformers.
    """

    name = "sentence_transformer"

    def __init__(
        self,
        *,
        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
        batch_size: int = 8,
    ) -> None:
        self._model_name = model_name
        self._batch_size = batch_size
        self._embedder: SentenceTransformerBasedEmbedder | None = None

    def _ensure_embedder(self) -> SentenceTransformerBasedEmbedder:
        if self._embedder is None:
            self._embedder = SentenceTransformerBasedEmbedder(
                model_name=self._model_name,
                batch_size=self._batch_size,
            )
        return self._embedder

    async def embed(self, texts: Sequence[str]) -> Sequence[Sequence[float]]:
        if not texts:
            return []

        embedder = self._ensure_embedder()

        def _run() -> List[List[float]]:
            vectors = embedder.embed(list(texts), batch_size=self._batch_size)
            return [vector.tolist() for vector in vectors]

        return await anyio.to_thread.run_sync(_run)


class OpenAIEmbeddingStrategy(EmbeddingStrategy):
    """
    Remote embedding backend using OpenAI's embeddings API.
    """

    name = "openai"

    def __init__(self, *, model_name: str, batch_size: int = 16) -> None:
        self._model_name = model_name
        self._batch_size = batch_size
        self._client: OpenAI | None = None

    def _client_instance(self) -> OpenAI:
        if self._client is None:
            self._client = OpenAI()
        return self._client

    async def embed(self, texts: Sequence[str]) -> Sequence[Sequence[float]]:
        if not texts:
            return []

        def _run() -> List[List[float]]:
            client = self._client_instance()
            vectors: List[List[float]] = []
            for start in range(0, len(texts), self._batch_size):
                batch = list(texts[start : start + self._batch_size])
                response: CreateEmbeddingResponse = client.embeddings.create(
                    model=self._model_name,
                    input=batch,
                )
                vectors.extend(data.embedding for data in response.data)
            return vectors

        return await anyio.to_thread.run_sync(_run)
