"""
OpenAI embeddings client implementation.
"""

import logging
import time
from typing import Any, Dict, List, Optional

from openai import AsyncOpenAI

from agentbyte.component import Component
from agentbyte.middleware.base import BaseMiddleware

from .._retry_observability import build_retry_summary
from ..base import AuthenticationError, InvalidRequestError, RateLimitError
from ..embeddings_base import BaseEmbeddingClient, BaseEmbeddingClientConfig
from ..pricing import PricingRegistry
from ..retry_policy import RetryMixin, RetryPolicy
from ..types import EmbeddingResult, ModelClientError, ModelPricing, Usage

logger = logging.getLogger(__name__)


class OpenAIEmbeddingClientConfig(BaseEmbeddingClientConfig):
    """Configuration for OpenAIEmbeddingClient serialization."""

    retry_policy: RetryPolicy = RetryPolicy()
    key_env_var: str = "OPENAI_API_KEY"


class OpenAIEmbeddingClient(
    BaseEmbeddingClient,
    Component[OpenAIEmbeddingClientConfig],
    RetryMixin,
):
    """OpenAI embeddings client with single and batch methods."""

    component_type = "model_client"
    component_config_schema = OpenAIEmbeddingClientConfig
    component_provider_override = "agentbyte.llm.OpenAIEmbeddingClient"

    def __init__(
        self,
        model: str,
        client: AsyncOpenAI,
        config: Optional[Dict[str, Any]] = None,
        middlewares: Optional[List[BaseMiddleware]] = None,
        model_pricing: Optional[ModelPricing] = None,
        max_retries: int = 3,
        initial_retry_delay: float = 1.0,
        max_retry_delay: float = 60.0,
        retry_policy: Optional[RetryPolicy] = None,
    ):
        super().__init__(model, client, config, middlewares=middlewares)
        self.model_pricing = model_pricing
        if retry_policy is not None:
            self.retry_policy = retry_policy
        else:
            self.retry_policy = RetryPolicy(
                max_retries=max_retries,
                initial_retry_delay=initial_retry_delay,
                max_retry_delay=max_retry_delay,
            )

    @property
    def max_retries(self) -> int:
        return self.retry_policy.max_retries

    @property
    def initial_retry_delay(self) -> float:
        return self.retry_policy.initial_retry_delay

    @property
    def max_retry_delay(self) -> float:
        return self.retry_policy.max_retry_delay

    @classmethod
    def from_api_key(
        cls,
        model: str,
        api_key: Optional[str] = None,
        org: Optional[str] = None,
        config: Optional[Dict[str, Any]] = None,
        middlewares: Optional[List[BaseMiddleware]] = None,
        model_pricing: Optional[ModelPricing] = None,
        max_retries: int = 3,
        initial_retry_delay: float = 1.0,
        max_retry_delay: float = 60.0,
        retry_policy: Optional[RetryPolicy] = None,
    ) -> "OpenAIEmbeddingClient":
        """Create an OpenAIEmbeddingClient from explicit or env-backed API key."""
        from pydantic import ValidationError

        from .settings import OpenAISettings

        if api_key is None:
            try:
                settings = OpenAISettings()
                resolved_key = settings.api_key.get_secret_value()
                resolved_org = org or settings.org
            except ValidationError as exc:
                raise ValueError(
                    "No OpenAI API key provided. Pass api_key= explicitly or set "
                    "the OPENAI_API_KEY environment variable."
                ) from exc
        else:
            resolved_key = api_key
            resolved_org = org

        openai_client = AsyncOpenAI(api_key=resolved_key, organization=resolved_org)
        return cls(
            model=model,
            client=openai_client,
            config=config,
            middlewares=middlewares,
            model_pricing=model_pricing,
            max_retries=max_retries,
            initial_retry_delay=initial_retry_delay,
            max_retry_delay=max_retry_delay,
            retry_policy=retry_policy,
        )

    async def create(self, input_text: str, **kwargs: Any) -> EmbeddingResult:
        """Create one embedding result from a single text input."""
        return await self.create_batch([input_text], **kwargs)

    async def create_batch(
        self,
        input_texts: List[str],
        **kwargs: Any,
    ) -> EmbeddingResult:
        """Create one embedding result from a batch of text inputs."""
        return await self._execute_embedding_operation(
            input_texts=input_texts,
            operation_kwargs=dict(kwargs),
            executor=self._create_batch_result_internal,
        )

    async def _create_batch_result_internal(
        self,
        input_texts: List[str],
        operation_kwargs: Dict[str, Any],
    ) -> EmbeddingResult:
        """Execute provider batch call and map response into EmbeddingResult."""
        try:
            retry_history: List = []
            start_time = time.perf_counter()
            response = await self._retry_with_backoff(
                self._create_batch_raw,
                input_texts,
                _retry_history=retry_history,
                **operation_kwargs,
            )
            duration_ms = int((time.perf_counter() - start_time) * 1000)

            vectors = self._extract_vectors(response, len(input_texts))
            prompt_tokens = (
                getattr(response.usage, "prompt_tokens", 0)
                if hasattr(response, "usage") and response.usage
                else 0
            )

            usage = Usage(
                duration_ms=duration_ms,
                llm_calls=1,
                tokens_input=prompt_tokens,
                tokens_output=0,
                cost_estimate=self._estimate_cost(
                    prompt_tokens=prompt_tokens,
                    completion_tokens=0,
                ),
            )

            metadata: Dict[str, Any] = {
                "input_count": len(input_texts),
                "total_tokens": (
                    getattr(response.usage, "total_tokens", None)
                    if hasattr(response, "usage") and response.usage
                    else None
                ),
            }
            if retry_history:
                metadata["retry_observability"] = {
                    "history": retry_history,
                    "latest": build_retry_summary(retry_history),
                }

            return EmbeddingResult(
                embedding=vectors[0],
                embeddings=vectors,
                model=getattr(response, "model", self.model),
                usage=usage,
                metadata=metadata,
            )
        except Exception as exc:
            self._handle_error(exc)

    async def _create_batch_raw(
        self,
        input_texts: List[str],
        **kwargs: Any,
    ) -> Any:
        """Internal raw API call for batch embeddings."""
        api_kwargs = {**self.config, **kwargs}
        return await self.client.embeddings.create(
            input=input_texts,
            model=self.model,
            **api_kwargs,
        )

    @staticmethod
    def _extract_vectors(response: Any, input_count: int) -> List[List[float]]:
        """Extract embedding vectors in request order from provider response."""
        vectors: List[List[float]] = [[] for _ in range(input_count)]
        fallback_vectors: List[List[float]] = []
        for item in response.data:
            embedding = list(item.embedding)
            index = getattr(item, "index", None)
            if isinstance(index, int) and 0 <= index < len(vectors):
                vectors[index] = embedding
            else:
                fallback_vectors.append(embedding)

        if fallback_vectors:
            return fallback_vectors
        return vectors

    def _estimate_cost(self, prompt_tokens: int, completion_tokens: int = 0) -> float:
        """Estimate embedding call cost from token usage."""
        if self.model_pricing is not None:
            pricing = self.model_pricing
        else:
            pricing = PricingRegistry.resolve_openai_embedding_pricing(self.model)
        return pricing.calculate_cost(prompt_tokens, completion_tokens)

    def _to_config(self) -> OpenAIEmbeddingClientConfig:
        """Serialize client to config without persisting secret values."""
        return OpenAIEmbeddingClientConfig(
            model=self.model,
            config=self.config.copy(),
            retry_policy=self.retry_policy,
            key_env_var=getattr(self, "_key_env_var", "OPENAI_API_KEY"),
        )

    @classmethod
    def _from_config(cls, config: OpenAIEmbeddingClientConfig) -> "OpenAIEmbeddingClient":
        """Reconstruct client from serialized config using env-resolved key."""
        import os

        api_key = os.environ.get(config.key_env_var)
        if not api_key:
            raise ValueError(
                f"Cannot reconstruct OpenAIEmbeddingClient: environment variable "
                f"'{config.key_env_var}' is not set."
            )
        instance = cls.from_api_key(
            model=config.model,
            api_key=api_key,
            config=config.config or None,
            retry_policy=config.retry_policy,
        )
        instance._key_env_var = config.key_env_var
        return instance

    def _handle_error(self, error: Exception) -> None:
        """Map OpenAI SDK exceptions to the shared error hierarchy."""
        try:
            from openai import (
                APIError,
                AuthenticationError as OpenAIAuthenticationError,
                BadRequestError,
                RateLimitError as OpenAIRateLimitError,
            )
        except ImportError:
            raise ModelClientError(f"OpenAI error: {str(error)}") from error

        error_msg = str(error)
        error_type = type(error).__name__

        if isinstance(error, OpenAIRateLimitError):
            raise RateLimitError(f"OpenAI rate limit exceeded: {error_msg}") from error
        if isinstance(error, OpenAIAuthenticationError):
            raise AuthenticationError(f"OpenAI authentication failed: {error_msg}") from error
        if isinstance(error, BadRequestError):
            raise InvalidRequestError(f"OpenAI request invalid: {error_msg}") from error
        if isinstance(error, APIError):
            raise ModelClientError(f"OpenAI API error ({error_type}): {error_msg}") from error
        raise ModelClientError(
            f"Unexpected error from OpenAI ({error_type}): {error_msg}"
        ) from error


__all__ = ["OpenAIEmbeddingClient", "OpenAIEmbeddingClientConfig"]
