"""
Azure OpenAI provider integration for Amplitude AI SDK.

Drop-in replacement that delegates to the ``openai.AzureOpenAI`` client
while preserving all Amplitude tracking from the standard OpenAI wrapper.

Usage::

    from amplitude import Amplitude
    from amplitude_ai.providers.azure_openai import AzureOpenAI

    amplitude = Amplitude("YOUR_AMP_KEY")
    client = AzureOpenAI(
        amplitude=amplitude,
        azure_endpoint="https://my-resource.openai.azure.com",
        api_key="AZURE_KEY",
        api_version="2024-06-01",
    )

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Hello!"}],
        amplitude_user_id="user-1",
    )
"""
from __future__ import annotations

import time
import uuid
from typing import Any, Dict, Iterator, List, Optional

try:
    from openai import AzureOpenAI as AzureOpenAIClient

    AZURE_OPENAI_AVAILABLE = True
except ImportError:
    AZURE_OPENAI_AVAILABLE = False
    AzureOpenAIClient = object  # type: ignore[assignment,misc]

from amplitude import Amplitude

from ..core.privacy import PrivacyConfig
from ..exceptions import ProviderError
from ..utils.logger import get_logger
from .base import BaseAIProvider, AmplitudeOrAI

# Re-use wrapped interfaces from the OpenAI provider
from .openai import WrappedChat, WrappedResponses


class AzureOpenAI(BaseAIProvider, AzureOpenAIClient):  # type: ignore[misc]
    """
    Azure OpenAI client wrapper with automatic Amplitude tracking.

    Inherits from :class:`BaseAIProvider` (tracking) and ``openai.AzureOpenAI``
    (API surface).  The tracking behaviour is identical to the standard
    :class:`amplitude_ai.providers.openai.OpenAI` wrapper — only the
    underlying client class differs.
    """

    def __init__(
        self,
        amplitude: AmplitudeOrAI,
        *,
        azure_endpoint: Optional[str] = None,
        api_key: Optional[str] = None,
        api_version: Optional[str] = None,
        privacy_mode: bool = False,
        content_mode: Optional[str] = None,
        **azure_kwargs: Any,
    ) -> None:
        """
        Initialise the Azure OpenAI client with Amplitude tracking.

        Args:
            amplitude: ``Amplitude`` client or ``AmplitudeAI`` instance.
            azure_endpoint: Azure resource endpoint URL.
            api_key: Azure OpenAI API key.
            api_version: Azure API version string.
            privacy_mode: Enable privacy mode (content hashing). Deprecated
                in favor of *content_mode*.
            content_mode: One of ``"full"``, ``"metadata_only"``,
                ``"customer_enriched"``. When provided, takes precedence
                over *privacy_mode*.
            **azure_kwargs: Additional arguments forwarded to
                ``openai.AzureOpenAI``.
        """
        if not AZURE_OPENAI_AVAILABLE:
            raise ProviderError(
                "OpenAI SDK is not installed or does not include AzureOpenAI. "
                "Install with: pip install openai"
            )

        privacy_config = PrivacyConfig(privacy_mode=privacy_mode, content_mode=content_mode)
        BaseAIProvider.__init__(
            self,
            amplitude=amplitude,
            privacy_config=privacy_config,
        )

        init_kwargs: Dict[str, Any] = {**azure_kwargs}
        if azure_endpoint is not None:
            init_kwargs["azure_endpoint"] = azure_endpoint
        if api_key is not None:
            init_kwargs["api_key"] = api_key
        if api_version is not None:
            init_kwargs["api_version"] = api_version

        AzureOpenAIClient.__init__(self, **init_kwargs)

        # Wrap the chat completions interface (same wrappers as standard OpenAI)
        self._original_chat = self.chat  # type: ignore[has-type]
        self.chat = WrappedChat(self)  # type: ignore[assignment]

        if hasattr(self, "responses"):
            self._original_responses = self.responses  # type: ignore[has-type]
            self.responses = WrappedResponses(self)  # type: ignore[assignment]

    # ------------------------------------------------------------------
    # BaseAIProvider abstract methods
    # ------------------------------------------------------------------

    def _get_provider_name(self) -> str:
        return "azure_openai"

    def _extract_model_name(self, **kwargs: Any) -> str:
        return kwargs.get("model", "unknown")

    def _extract_messages(self, **kwargs: Any) -> List[Dict[str, Any]]:
        return kwargs.get("messages", [])

    def _format_messages_for_tracking(
        self, messages: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        from ..core.privacy import extract_text_from_openai_message, sanitize_any_content

        formatted = []
        for msg in messages:
            formatted_msg = msg.copy()
            text_content = extract_text_from_openai_message(msg)
            if text_content:
                sanitized = sanitize_any_content(
                    text_content,
                    privacy_mode=self.privacy_config.privacy_mode,
                    redact_pii=self.privacy_config.redact_pii,
                )
                if self.privacy_config.privacy_mode:
                    formatted_msg["content_hash"] = sanitized.get("content_hash")
                    formatted_msg.pop("content", None)
                else:
                    formatted_msg["content"] = sanitized["$llm_message"]["text"]
            formatted.append(formatted_msg)
        return formatted

    def _extract_response_content(self, response: Any) -> str:
        if hasattr(response, "choices") and response.choices:
            choice = response.choices[0]
            if hasattr(choice, "message") and hasattr(choice.message, "content"):
                return choice.message.content or ""
        return ""

    def _extract_finish_reason(self, response: Any) -> Optional[str]:
        if hasattr(response, "choices") and response.choices:
            return getattr(response.choices[0], "finish_reason", None)
        return None

    def _extract_reasoning_content(self, response: Any) -> Optional[str]:
        if hasattr(response, "choices") and response.choices:
            message = getattr(response.choices[0], "message", None)
            if message is not None:
                reasoning = getattr(message, "reasoning_content", None)
                if reasoning:
                    return reasoning  # type: ignore[no-any-return]
        return None

    def _extract_tool_calls(self, response: Any) -> Optional[List[Dict[str, Any]]]:
        if hasattr(response, "choices") and response.choices:
            choice = response.choices[0]
            if hasattr(choice, "message") and hasattr(choice.message, "tool_calls"):
                tool_calls = choice.message.tool_calls
                if tool_calls:
                    return [
                        {
                            "type": tc.type,
                            "id": tc.id,
                            "function": {
                                "name": tc.function.name,
                                "arguments": tc.function.arguments,
                            },
                        }
                        for tc in tool_calls
                    ]
        return None

    def __getattr__(self, name: str) -> Any:
        if name.startswith("_") or name in (
            "amplitude",
            "privacy_config",
            "provider_name",
        ):
            raise AttributeError(
                f"'{type(self).__name__}' object has no attribute '{name}'"
            )
        return getattr(super(AzureOpenAIClient, self), name)  # type: ignore[arg-type]
