"""OpenAI embedding provider."""

from __future__ import annotations

import os
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from openai import AsyncOpenAI

# Default dimensions per model
_MODEL_DIMENSIONS: dict[str, int] = {
    "text-embedding-3-small": 1536,
    "text-embedding-3-large": 3072,
    "text-embedding-ada-002": 1536,
}

# Valid dimension reductions for embedding-3 models
_VALID_DIMENSIONS = {256, 512, 768, 1024, 1536, 2048, 3072}


class OpenAIEmbedding:
    """OpenAI embedding provider (text-embedding-3-small, text-embedding-3-large).

    Supports dimension reduction for text-embedding-3 models.
    Uses OPENAI_API_KEY environment variable if api_key not provided.

    Example:
        provider = OpenAIEmbedding()
        embeddings = await provider.embed(["hello world"])

        # With dimension reduction
        provider = OpenAIEmbedding(dimensions=256)
        embeddings = await provider.embed(["hello world"])  # 256-dim vectors
    """

    def __init__(
        self,
        model: str = "text-embedding-3-small",
        api_key: str | None = None,
        dimensions: int | None = None,
        batch_size: int = 100,
    ) -> None:
        """Initialize OpenAI embedding provider.

        Args:
            model: Model name. Defaults to text-embedding-3-small.
            api_key: OpenAI API key. Defaults to OPENAI_API_KEY env var.
            dimensions: Optional dimension reduction (256, 512, 768, 1024, 1536, 2048, 3072).
                       Only supported for text-embedding-3 models.
            batch_size: Max texts per API call. P8: defaults to 100 to avoid rate-limit
                       and payload-size issues with large corpora.
        """
        self._model = model
        self._api_key = api_key or os.getenv("OPENAI_API_KEY")
        self._batch_size = batch_size

        if not self._api_key:
            raise ValueError(
                "API key required. Set OPENAI_API_KEY environment variable "
                "or pass api_key explicitly."
            )

        # Validate dimensions
        if dimensions is not None:
            if dimensions not in _VALID_DIMENSIONS:
                raise ValueError(
                    f"dimensions must be one of {_VALID_DIMENSIONS} "
                    f"for text-embedding-3 models, got {dimensions}"
                )
            if model not in ("text-embedding-3-small", "text-embedding-3-large"):
                raise ValueError(
                    f"dimension reduction only supported for text-embedding-3 models, got {model}"
                )
            self._dimensions = dimensions
        else:
            self._dimensions = _MODEL_DIMENSIONS.get(model, 1536)

        self._client: AsyncOpenAI | None = None

    @property
    def batch_size(self) -> int:
        """Max texts per API call."""
        return self._batch_size

    @property
    def dimensions(self) -> int:
        """Vector dimensionality."""
        return self._dimensions

    @property
    def model(self) -> str:
        """Model name."""
        return self._model

    @property
    def model_id(self) -> str:
        """Model identifier for cost tracking."""
        return self._model

    def _get_client(self) -> AsyncOpenAI:
        """Get or create async OpenAI client."""
        if self._client is None:
            from openai import AsyncOpenAI

            self._client = AsyncOpenAI(api_key=self._api_key)
        return self._client

    async def embed(
        self,
        texts: list[str],
        budget_tracker: object | None = None,
    ) -> list[list[float]]:
        """Embed texts into vectors.

        Args:
            texts: List of text strings to embed.
            budget_tracker: Optional BudgetTracker for cost tracking.

        Returns:
            List of embedding vectors.
        """
        if not texts:
            return []

        client = self._get_client()
        all_embeddings: list[list[float]] = []

        # P8: split into batches to avoid payload-size and rate-limit issues
        for i in range(0, len(texts), self._batch_size):
            batch = texts[i : i + self._batch_size]

            kwargs: dict[str, object] = {
                "model": self._model,
                "input": batch,
            }

            # Add dimensions if reduced
            if self._dimensions != _MODEL_DIMENSIONS.get(self._model, 1536):
                kwargs["dimensions"] = self._dimensions

            response = await client.embeddings.create(**kwargs)

            all_embeddings.extend(item.embedding for item in response.data)

            # Track cost if budget_tracker provided
            if budget_tracker is not None and response.usage is not None:
                from syrin.cost import calculate_embedding_cost

                token_count = response.usage.prompt_tokens
                cost = calculate_embedding_cost(self._model, token_count)
                budget_tracker.record_external(  # type: ignore[attr-defined]
                    service="openai_embedding",
                    cost_usd=cost,
                    metadata={
                        "model": self._model,
                        "token_count": token_count,
                    },
                )

        return all_embeddings
