"""
Azure OpenAI chat completion client implementation.
"""

import json
import logging
import time
from collections.abc import AsyncGenerator
from typing import Any, Dict, List, Optional, TypeVar

from openai import AsyncAzureOpenAI

from agentbyte.messages import AssistantMessage, Message, ToolCallRequest

from agentbyte.component import Component

from .._retry_observability import build_retry_summary
from ..base import (
    BaseChatCompletionClient,
    BaseChatCompletionClientConfig,
    RateLimitError,
    AuthenticationError,
    InvalidRequestError,
)
from ..pricing import PricingRegistry
from ..retry_policy import RetryMixin, RetryPolicy
from ..types import (
    ChatCompletionChunk,
    ChatCompletionResult,
    ModelPricing,
    ModelClientError,
    Usage,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


class AzureOpenAIChatCompletionClientConfig(BaseChatCompletionClientConfig):
    """Configuration for AzureOpenAIChatCompletionClient serialization."""

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


class AzureOpenAIChatCompletionClient(
    BaseChatCompletionClient,
    Component[AzureOpenAIChatCompletionClientConfig],
    RetryMixin,
):
    """Azure OpenAI chat completion client."""

    component_type = "model_client"
    component_config_schema = AzureOpenAIChatCompletionClientConfig
    component_provider_override = "agentbyte.llm.AzureOpenAIChatCompletionClient"

    def __init__(
        self,
        model: str,
        client: AsyncAzureOpenAI,
        config: Optional[Dict[str, Any]] = 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)
        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,
        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,
    ) -> "AzureOpenAIChatCompletionClient":
        """Create a client using an Azure OpenAI API key."""
        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.chat_deployment_name
        if not resolved_model:
            raise ValueError(
                "No deployment name provided. Pass model= explicitly or set "
                "AZURE_OPENAI_CHAT_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,
            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,
        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,
    ) -> "AzureOpenAIChatCompletionClient":
        """Create a client using DefaultAzureCredential."""
        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.chat_deployment_name
        if not resolved_model:
            raise ValueError(
                "No deployment name provided. Pass model= explicitly or set "
                "AZURE_OPENAI_CHAT_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,
            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,
        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,
    ) -> "AzureOpenAIChatCompletionClient":
        """Create a client using certificate-based Azure AD service principal auth."""
        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.chat_deployment_name
        if not resolved_model:
            raise ValueError(
                "No deployment name provided. Pass model= explicitly or set "
                "AZURE_OPENAI_CHAT_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,
            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,
        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,
    ) -> "AzureOpenAIChatCompletionClient":
        """Create a 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.chat_deployment_name
        if not resolved_model:
            raise ValueError(
                "No deployment name provided. Pass model= explicitly or set "
                "AZURE_OPENAI_CHAT_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,
            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,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        """Make a single Azure OpenAI API call."""
        retry_history: List = []
        try:
            result = await self._retry_with_backoff(
                self._create_internal,
                messages,
                tools,
                output_format,
                _retry_history=retry_history,
                **kwargs,
            )
            if retry_history:
                result.metadata["retry_observability"] = {
                    "history": retry_history,
                    "latest": build_retry_summary(retry_history),
                }
            return result
        except Exception as e:
            self._handle_error(e)

    async def _create_internal(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        """Internal create implementation (retried by wrapper)."""
        api_kwargs = {**self.config, **kwargs}
        api_messages = self._convert_messages_to_api_format(messages)

        structured_output = None
        if output_format:
            try:
                schema = output_format.model_json_schema()
                schema = self._make_schema_compatible(schema)
                api_kwargs["response_format"] = {
                    "type": "json_schema",
                    "json_schema": {
                        "name": schema.get("title", output_format.__name__),
                        "description": schema.get(
                            "description",
                            f"Structured output for {output_format.__name__}",
                        ),
                        "strict": True,
                        "schema": schema,
                    },
                }
            except Exception as e:
                logger.warning(
                    f"Failed to convert {output_format.__name__} to JSON schema: {e}. "
                    f"Continuing without structured output."
                )

        start_time = time.perf_counter()
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=api_messages,
            tools=tools,
            **api_kwargs,
        )
        duration_ms = int((time.perf_counter() - start_time) * 1000)

        choice = response.choices[0]
        message_content = choice.message.content or ""

        tool_calls = None
        if choice.message.tool_calls:
            tool_calls = [
                ToolCallRequest(
                    call_id=tc.id,
                    tool_name=tc.function.name,
                    parameters=json.loads(tc.function.arguments),
                )
                for tc in choice.message.tool_calls
            ]

        if output_format and message_content:
            try:
                structured_output = output_format.model_validate_json(message_content)
            except Exception as e:
                logger.warning(f"Failed to parse structured output: {e}. Using raw content.")
                structured_output = None

        return ChatCompletionResult(
            message=AssistantMessage(
                content=message_content,
                source=self.model,
                tool_calls=tool_calls,
            ),
            usage=Usage(
                duration_ms=duration_ms,
                llm_calls=1,
                tokens_input=response.usage.prompt_tokens,
                tokens_output=response.usage.completion_tokens,
                tool_calls=len(tool_calls) if tool_calls else 0,
                cost_estimate=self._estimate_cost(
                    response.usage.prompt_tokens,
                    response.usage.completion_tokens,
                ),
            ),
            model=response.model,
            metadata={"finish_reason": choice.finish_reason},
            finish_reason=choice.finish_reason,
            structured_output=structured_output,
        )

    async def create_stream(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        stream_options: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        """Make a streaming Azure OpenAI API call."""
        try:
            async for chunk in self._retry_stream_with_backoff(
                self._create_stream_internal,
                messages,
                tools,
                output_format,
                stream_options,
                **kwargs,
            ):
                yield chunk
        except Exception as e:
            self._handle_error(e)

    async def _create_stream_internal(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        stream_options: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        """Internal streaming implementation (retried by wrapper)."""
        api_kwargs = {**self.config, **kwargs}
        if "stream_options" in api_kwargs and stream_options is None:
            stream_options = api_kwargs.pop("stream_options")
        else:
            api_kwargs.pop("stream_options", None)
        if stream_options is None:
            stream_options = {"include_usage": True}
        if stream_options:
            api_kwargs["stream_options"] = stream_options

        api_messages = self._convert_messages_to_api_format(messages)

        start_time = time.perf_counter()
        stream = await self.client.chat.completions.create(
            model=self.model,
            messages=api_messages,
            tools=tools,
            stream=True,
            **api_kwargs,
        )

        tool_call_chunks: Dict[Any, Dict[str, Any]] = {}

        async for chunk in stream:
            if hasattr(chunk, "usage") and chunk.usage and (not chunk.choices or len(chunk.choices) == 0):
                duration_ms = int((time.perf_counter() - start_time) * 1000)
                usage = Usage(
                    duration_ms=duration_ms,
                    llm_calls=1,
                    tokens_input=chunk.usage.prompt_tokens,
                    tokens_output=chunk.usage.completion_tokens,
                    cost_estimate=self._estimate_cost(
                        chunk.usage.prompt_tokens,
                        chunk.usage.completion_tokens,
                    ),
                )
                yield ChatCompletionChunk(
                    content="",
                    model=chunk.model,
                    is_complete=True,
                    tool_call_chunk=None,
                    usage=usage,
                    metadata={},
                )
                break

            if not chunk.choices:
                continue

            chunk_choice = chunk.choices[0]
            delta = chunk_choice.delta

            if delta.content:
                yield ChatCompletionChunk(
                    content=delta.content,
                    model=chunk.model,
                    is_complete=False,
                    tool_call_chunk=None,
                    usage=None,
                    metadata={"finish_reason": chunk_choice.finish_reason},
                )

            if delta.tool_calls:
                for tc_delta in delta.tool_calls:
                    index = getattr(tc_delta, "index", None)
                    call_id = tc_delta.id
                    tracking_key = index if index is not None else call_id

                    if tracking_key not in tool_call_chunks:
                        tool_call_chunks[tracking_key] = {
                            "id": call_id,
                            "function": {"name": "", "arguments": ""},
                        }

                    if call_id:
                        tool_call_chunks[tracking_key]["id"] = call_id

                    if tc_delta.function:
                        if tc_delta.function.name:
                            tool_call_chunks[tracking_key]["function"]["name"] = tc_delta.function.name
                        if tc_delta.function.arguments:
                            tool_call_chunks[tracking_key]["function"]["arguments"] += tc_delta.function.arguments

                    yield ChatCompletionChunk(
                        content="",
                        model=chunk.model,
                        is_complete=False,
                        tool_call_chunk=tool_call_chunks[tracking_key],
                        usage=None,
                        metadata={},
                    )

            if chunk_choice.finish_reason and not stream_options.get("include_usage"):
                yield ChatCompletionChunk(
                    content="",
                    model=chunk.model,
                    is_complete=True,
                    tool_call_chunk=None,
                    usage=None,
                    metadata={"finish_reason": chunk_choice.finish_reason},
                )

    def _make_schema_compatible(self, schema: Dict[str, Any]) -> Dict[str, Any]:
        """Make JSON schema compatible with Azure OpenAI Structured Outputs."""
        compatible_schema: Dict[str, Any] = {}
        for key, value in schema.items():
            if isinstance(value, dict):
                compatible_schema[key] = self._make_schema_compatible(value)
            elif isinstance(value, list):
                compatible_schema[key] = [
                    self._make_schema_compatible(item) if isinstance(item, dict) else item
                    for item in value
                ]
            else:
                compatible_schema[key] = value

        if compatible_schema.get("type") == "object":
            compatible_schema["additionalProperties"] = False

        return compatible_schema

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

    def _to_config(self) -> AzureOpenAIChatCompletionClientConfig:
        """Serialize to config without persisting secrets."""
        return AzureOpenAIChatCompletionClientConfig(
            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_CHAT_DEPLOYMENT_NAME"),
            api_version_env_var=getattr(self, "_api_version_env_var", "AZURE_OPENAI_API_VERSION"),
        )

    @classmethod
    def _from_config(
        cls, config: AzureOpenAIChatCompletionClientConfig
    ) -> "AzureOpenAIChatCompletionClient":
        """Reconstruct client from saved config."""
        import os

        resolved_key = os.environ.get(config.key_env_var)
        if not resolved_key:
            raise ValueError(
                f"Cannot reconstruct AzureOpenAIChatCompletionClient: environment variable "
                f"'{config.key_env_var}' is not set."
            )
        instance = cls.from_api_key(
            model=config.model if config.model != "gpt-4.1-mini" else None,
            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:
        """Convert Azure OpenAI exceptions to unified error hierarchy."""
        try:
            from openai import (
                RateLimitError as OpenAIRateLimitError,
                AuthenticationError as OpenAIAuthenticationError,
                BadRequestError,
                APIError,
            )
        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
        elif isinstance(error, OpenAIAuthenticationError):
            raise AuthenticationError(
                f"Azure OpenAI authentication failed: {error_msg}"
            ) from error
        elif isinstance(error, BadRequestError):
            raise InvalidRequestError(f"Azure OpenAI request invalid: {error_msg}") from error
        elif isinstance(error, APIError):
            raise ModelClientError(
                f"Azure OpenAI API error ({error_type}): {error_msg}"
            ) from error
        else:
            raise ModelClientError(
                f"Unexpected error from Azure OpenAI ({error_type}): {error_msg}"
            ) from error


__all__ = ["AzureOpenAIChatCompletionClient", "AzureOpenAIChatCompletionClientConfig"]
