"""OpenAI embeddings client (also works against OpenAI-compatible endpoints)."""

from __future__ import annotations

import httpx

from legate.rag.embeddings.base import EmbeddingError


class OpenAIEmbedder:
    provider = "openai"

    def __init__(
        self,
        *,
        api_key: str,
        model: str = "text-embedding-3-small",
        base_url: str = "https://api.openai.com/v1",
        timeout: float = 60.0,
        http_client: httpx.Client | None = None,
    ) -> None:
        self._api_key = api_key
        self.model = model
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._http = http_client

    def embed(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        payload = {"model": self.model, "input": texts}
        headers = {"Authorization": f"Bearer {self._api_key}"}
        try:
            if self._http is not None:
                resp = self._http.post(
                    f"{self._base_url}/embeddings", json=payload, headers=headers
                )
            else:
                with httpx.Client(timeout=self._timeout) as client:
                    resp = client.post(
                        f"{self._base_url}/embeddings", json=payload, headers=headers
                    )
            resp.raise_for_status()
            data = resp.json()
        except httpx.HTTPError as exc:
            raise EmbeddingError(f"openai embeddings request failed: {exc}") from exc

        try:
            items = sorted(data["data"], key=lambda d: d["index"])
            return [item["embedding"] for item in items]
        except (KeyError, TypeError) as exc:
            raise EmbeddingError(f"openai returned an unexpected response: {data}") from exc
