"""
Azure OpenAI embeddings client implementation.
"""

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

from openai import AsyncAzureOpenAI

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 AzureOpenAIEmbeddingClientConfig(BaseEmbeddingClientConfig):
    """Configuration for AzureOpenAIEmbeddingClient serialization."""

    retry_policy: RetryPolicy = RetryPolicy()
    key_env_var: str = "AZURE_OPENAI_API_KEY"
    deployment_env_var: str = "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"
    api_version_env_var: str = "AZURE_OPENAI_API_VERSION"


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

    component_type = "model_client"
    component_config_schema = AzureOpenAIEmbeddingClientConfig
    component_provider_override = "agentbyte.llm.AzureOpenAIEmbeddingClient"

    def __init__(
        self,
        model: str,
        client: AsyncAzureOpenAI,
        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: Optional[str] = None,
        api_version: Optional[str] = None,
        api_key: Optional[str] = None,
        service_settings: Optional[Any] = 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,
    ) -> "AzureOpenAIEmbeddingClient":
        """Create client using Azure OpenAI API key authentication."""
        from pydantic import ValidationError

        from .settings import AzureOpenAISettings

        try:
            svc = service_settings or AzureOpenAISettings()
        except ValidationError as exc:
            raise ValueError(
                "Azure OpenAI service settings incomplete. Ensure AZURE_OPENAI_ENDPOINT "
                "is set, or pass service_settings= explicitly."
            ) from exc

        resolved_model = model or svc.embedding_deployment_name or svc.chat_deployment_name
        if not resolved_model:
            raise ValueError(
                "No embedding deployment name provided. Pass model= explicitly or set "
                "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME."
            )

        resolved_version = api_version or svc.api_version
        if not resolved_version:
            raise ValueError(
                "No API version provided. Pass api_version= explicitly or set "
                "AZURE_OPENAI_API_VERSION."
            )

        resolved_key = api_key or (svc.api_key.get_secret_value() if svc.api_key else None)
        if not resolved_key:
            raise ValueError(
                "No Azure OpenAI API key provided. Pass api_key= explicitly, set "
                "AZURE_OPENAI_API_KEY, or use a different auth method."
            )

        azure_client = AsyncAzureOpenAI(
            api_key=resolved_key,
            azure_endpoint=svc.endpoint,
            api_version=resolved_version,
        )
        return cls(
            model=resolved_model,
            client=azure_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,
        )

    @classmethod
    def from_default_credential(
        cls,
        model: Optional[str] = None,
        api_version: Optional[str] = None,
        service_settings: Optional[Any] = None,
        scope: str = "https://cognitiveservices.azure.com/.default",
        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,
    ) -> "AzureOpenAIEmbeddingClient":
        """Create client using DefaultAzureCredential token provider."""
        try:
            from azure.identity import DefaultAzureCredential, get_bearer_token_provider
        except ImportError as exc:
            raise ImportError(
                "azure-identity is required for from_default_credential(). "
                "Install it with: pip install azure-identity"
            ) from exc

        from pydantic import ValidationError

        from .settings import AzureOpenAISettings

        try:
            svc = service_settings or AzureOpenAISettings()
        except ValidationError as exc:
            raise ValueError(
                "Azure OpenAI service settings incomplete. Ensure AZURE_OPENAI_ENDPOINT "
                "is set, or pass service_settings= explicitly."
            ) from exc

        resolved_model = model or svc.embedding_deployment_name or svc.chat_deployment_name
        if not resolved_model:
            raise ValueError(
                "No embedding deployment name provided. Pass model= explicitly or set "
                "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME."
            )

        resolved_version = api_version or svc.api_version
        if not resolved_version:
            raise ValueError(
                "No API version provided. Pass api_version= explicitly or set "
                "AZURE_OPENAI_API_VERSION."
            )

        token_provider = get_bearer_token_provider(DefaultAzureCredential(), scope)
        azure_client = AsyncAzureOpenAI(
            azure_ad_token_provider=token_provider,
            azure_endpoint=svc.endpoint,
            api_version=resolved_version,
        )
        return cls(
            model=resolved_model,
            client=azure_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,
        )

    @classmethod
    def from_certificate(
        cls,
        model: Optional[str] = None,
        api_version: Optional[str] = None,
        service_settings: Optional[Any] = None,
        principal_settings: Optional[Any] = None,
        scope: 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,
    ) -> "AzureOpenAIEmbeddingClient":
        """Create client using certificate-based Azure AD authentication."""
        from pydantic import ValidationError

        from .auth import get_certificate_token_provider
        from .settings import AzureOpenAISettings, AzureServicePrincipalSettings

        try:
            svc = service_settings or AzureOpenAISettings()
        except ValidationError as exc:
            raise ValueError(
                "Azure OpenAI service settings incomplete. Ensure AZURE_OPENAI_ENDPOINT "
                "is set, or pass service_settings= explicitly."
            ) from exc

        try:
            sp = principal_settings or AzureServicePrincipalSettings()
        except ValidationError as exc:
            raise ValueError(
                "Azure service principal settings incomplete. Ensure AZURE_TENANT_ID, "
                "AZURE_CLIENT_ID, AZURE_COGNITIVE_SCOPE, AZURE_CLIENT_CERTIFICATE_SECRET, "
                "and AZURE_CLIENT_PRIVATE_KEY_SECRET are set, or pass principal_settings= explicitly."
            ) from exc

        resolved_model = model or svc.embedding_deployment_name or svc.chat_deployment_name
        if not resolved_model:
            raise ValueError(
                "No embedding deployment name provided. Pass model= explicitly or set "
                "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME."
            )

        resolved_version = api_version or svc.api_version
        if not resolved_version:
            raise ValueError(
                "No API version provided. Pass api_version= explicitly or set "
                "AZURE_OPENAI_API_VERSION."
            )

        resolved_scope = scope or sp.cognitive_scope
        token_provider = get_certificate_token_provider(
            tenant_id=sp.tenant_id,
            client_id=sp.client_id,
            certificate_data=sp.certificate_data,
            scope=resolved_scope,
        )
        azure_client = AsyncAzureOpenAI(
            azure_ad_token_provider=token_provider,
            azure_endpoint=svc.endpoint,
            api_version=resolved_version,
        )
        return cls(
            model=resolved_model,
            client=azure_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,
        )

    @classmethod
    def from_ad_token(
        cls,
        ad_token: str,
        model: Optional[str] = None,
        api_version: Optional[str] = None,
        service_settings: Optional[Any] = 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,
    ) -> "AzureOpenAIEmbeddingClient":
        """Create client using a pre-fetched Azure AD bearer token."""
        from pydantic import ValidationError

        from .settings import AzureOpenAISettings

        try:
            svc = service_settings or AzureOpenAISettings()
        except ValidationError as exc:
            raise ValueError(
                "Azure OpenAI service settings incomplete. Ensure AZURE_OPENAI_ENDPOINT "
                "is set, or pass service_settings= explicitly."
            ) from exc

        resolved_model = model or svc.embedding_deployment_name or svc.chat_deployment_name
        if not resolved_model:
            raise ValueError(
                "No embedding deployment name provided. Pass model= explicitly or set "
                "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME."
            )

        resolved_version = api_version or svc.api_version
        if not resolved_version:
            raise ValueError(
                "No API version provided. Pass api_version= explicitly or set "
                "AZURE_OPENAI_API_VERSION."
            )

        azure_client = AsyncAzureOpenAI(
            azure_ad_token=ad_token,
            azure_endpoint=svc.endpoint,
            api_version=resolved_version,
        )
        return cls(
            model=resolved_model,
            client=azure_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_azure_embedding_pricing(self.model)
        return pricing.calculate_cost(prompt_tokens, completion_tokens)

    def _to_config(self) -> AzureOpenAIEmbeddingClientConfig:
        """Serialize client to config without persisting secret values."""
        return AzureOpenAIEmbeddingClientConfig(
            model=self.model,
            config=self.config.copy(),
            retry_policy=self.retry_policy,
            key_env_var=getattr(self, "_key_env_var", "AZURE_OPENAI_API_KEY"),
            deployment_env_var=getattr(
                self, "_deployment_env_var", "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"
            ),
            api_version_env_var=getattr(self, "_api_version_env_var", "AZURE_OPENAI_API_VERSION"),
        )

    @classmethod
    def _from_config(
        cls,
        config: AzureOpenAIEmbeddingClientConfig,
    ) -> "AzureOpenAIEmbeddingClient":
        """Reconstruct client from serialized config using env-resolved settings."""
        import os

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

    def _handle_error(self, error: Exception) -> None:
        """Map Azure 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"Azure OpenAI error: {str(error)}") from error

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

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


__all__ = ["AzureOpenAIEmbeddingClient", "AzureOpenAIEmbeddingClientConfig"]
