# chuk_llm/llm/providers/azure_openai_client.py
"""
Azure OpenAI chat-completion adapter with unified configuration integration
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Enhanced wrapper around the official `openai` SDK configured for Azure OpenAI
that uses the unified configuration system for all capabilities.

COMPLETE VERSION WITH SMART DISCOVERY SUPPORT:
- Supports ANY Azure deployment name (no validation against static list)
- Smart defaults for discovered deployments
- Pattern-based capability detection
- Full compatibility with custom deployment names like "scribeflowgpt4o"
- Universal tool name compatibility with bidirectional mapping
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import uuid
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Any

import openai

if TYPE_CHECKING:
    from chuk_llm.core.models import Message, Tool

from chuk_llm.core.constants import AzureOpenAIParam
from chuk_llm.llm.providers._config_mixin import ConfigAwareProviderMixin

# mixins
from chuk_llm.llm.providers._mixins import OpenAIStyleMixin
from chuk_llm.llm.providers._tool_compatibility import ToolCompatibilityMixin

# base
from ..core.base import BaseLLMClient

log = logging.getLogger(__name__)


class AzureOpenAILLMClient(
    ConfigAwareProviderMixin, ToolCompatibilityMixin, OpenAIStyleMixin, BaseLLMClient
):
    """
    Configuration-driven wrapper around the official `openai` SDK for Azure OpenAI
    that gets all capabilities from the unified YAML configuration.

    ENHANCED: Now supports discovered deployments with smart defaults!
    - ANY deployment name is valid (no static list validation)
    - Smart capability detection based on deployment name patterns
    - Automatic feature inference for unknown deployments

    Uses universal tool name compatibility system to handle any naming convention:
    - stdio.read_query -> stdio_read_query (if needed)
    - web.api:search -> web_api_search (if needed)
    - database.sql.execute -> database_sql_execute (if needed)
    """

    def __init__(
        self,
        model: str = "gpt-4o-mini",
        api_key: str | None = None,
        azure_endpoint: str | None = None,
        api_version: str | None = None,
        azure_deployment: str | None = None,
        azure_ad_token: str | None = None,
        azure_ad_token_provider: Any | None = None,
    ) -> None:
        # Initialize mixins
        ConfigAwareProviderMixin.__init__(self, "azure_openai", model)
        ToolCompatibilityMixin.__init__(self, "azure_openai")

        self.model = model
        self.azure_endpoint = azure_endpoint
        self.api_version = api_version or "2024-02-01"
        self.azure_deployment = (
            azure_deployment or model
        )  # Default deployment name to model name

        # Azure OpenAI client configuration
        client_kwargs = {
            AzureOpenAIParam.API_VERSION.value: self.api_version,
            AzureOpenAIParam.AZURE_ENDPOINT.value: azure_endpoint
            or os.getenv("AZURE_OPENAI_ENDPOINT"),
        }

        # Authentication - priority order: token provider > token > api key
        if azure_ad_token_provider:
            client_kwargs["azure_ad_token_provider"] = azure_ad_token_provider
        elif azure_ad_token:
            client_kwargs["azure_ad_token"] = azure_ad_token
        else:
            client_kwargs[AzureOpenAIParam.API_KEY.value] = api_key or os.getenv(
                "AZURE_OPENAI_API_KEY"
            )

        # Validate required parameters
        if not client_kwargs.get(AzureOpenAIParam.AZURE_ENDPOINT.value):
            raise ValueError(
                "azure_endpoint is required for Azure OpenAI. Set AZURE_OPENAI_ENDPOINT or pass azure_endpoint parameter."
            )

        if not any(
            [
                azure_ad_token_provider,
                azure_ad_token,
                client_kwargs.get(AzureOpenAIParam.API_KEY.value),
            ]
        ):
            raise ValueError(
                "Authentication required: provide api_key, azure_ad_token, or azure_ad_token_provider"
            )

        # Use AzureOpenAI for real streaming support
        self.async_client = openai.AsyncAzureOpenAI(**client_kwargs)  # type: ignore[call-overload]

        # Keep sync client for backwards compatibility if needed
        self.client = openai.AzureOpenAI(**client_kwargs)  # type: ignore[call-overload]

        log.debug(
            f"Azure OpenAI client initialized: endpoint={azure_endpoint}, deployment={self.azure_deployment}, model={self.model}"
        )

    # ================================================================
    # SMART DEFAULTS FOR DISCOVERED AZURE DEPLOYMENTS
    # ================================================================

    @staticmethod
    def _get_smart_default_features(deployment_name: str) -> set[str]:
        """
        Get smart default features for an Azure deployment based on naming patterns.

        This allows discovered deployments to work even if not in static config.
        For example, "scribeflowgpt4o" will be detected as a GPT-4 variant with full capabilities.
        """
        deployment_lower = deployment_name.lower()

        # Base features that ALL modern Azure OpenAI deployments should have
        base_features = {"text", "streaming", "system_messages"}

        # Pattern-based feature detection for Azure deployments
        if any(pattern in deployment_lower for pattern in ["o1", "o3", "o4", "o5"]):
            # Reasoning model deployments
            if "o1" in deployment_lower:
                # O1 models are legacy - no tools, limited capabilities
                return {"text", "reasoning"}
            else:
                # O3+ models are modern reasoning - assume full tool support
                return {"text", "streaming", "tools", "reasoning", "system_messages"}

        elif any(
            pattern in deployment_lower
            for pattern in ["gpt-4", "gpt4", "gpt-5", "gpt5", "gpt-3", "gpt3"]
        ):
            # Standard GPT model deployments - assume modern capabilities
            features = base_features | {"tools", "json_mode"}

            # Vision support for GPT-4+ models (not 3.5)
            if any(v in deployment_lower for v in ["gpt-4", "gpt4", "gpt-5", "gpt5"]):
                features.add("vision")
                features.add("parallel_calls")

            # GPT-5 models use reasoning architecture
            if any(v in deployment_lower for v in ["gpt-5", "gpt5"]):
                features.add("reasoning")

            return features

        elif "embedding" in deployment_lower or "ada" in deployment_lower:
            # Embedding models don't support chat features
            return {"text"}

        elif "whisper" in deployment_lower:
            # Audio models
            return {"audio", "transcription"}

        elif "dall-e" in deployment_lower or "dalle" in deployment_lower:
            # Image generation models
            return {"image_generation"}

        else:
            # Unknown Azure deployment patterns - be optimistic about capabilities
            # Azure deployments of unknown models likely support standard features
            log.info(
                f"Unknown Azure deployment pattern '{deployment_name}' - assuming modern chat capabilities including tools"
            )
            return base_features | {"tools", "json_mode"}

    @staticmethod
    def _get_smart_default_parameters(deployment_name: str) -> dict[str, Any]:
        """Get smart default parameters for an Azure deployment"""
        deployment_lower = deployment_name.lower()

        # Reasoning model parameter handling
        if any(
            pattern in deployment_lower
            for pattern in ["o1", "o3", "o4", "o5", "gpt-5", "gpt5"]
        ):
            return {
                "max_context_length": (
                    272000
                    if "gpt-5" in deployment_lower or "gpt5" in deployment_lower
                    else 200000
                ),
                "max_output_tokens": (
                    16384
                    if "gpt-5" in deployment_lower or "gpt5" in deployment_lower
                    else 65536
                ),
                "requires_max_completion_tokens": True,
                "parameter_mapping": {"max_tokens": "max_completion_tokens"},
                "unsupported_params": [
                    "temperature",
                    "top_p",
                    "frequency_penalty",
                    "presence_penalty",
                ],
                "supports_tools": "o1" not in deployment_lower,
            }

        # Standard model defaults - generous assumptions for new deployments
        return {
            "max_context_length": 128000,
            "max_output_tokens": 4096,
            "supports_tools": True,  # Assume new deployments support tools
        }

    def _has_explicit_deployment_config(self, deployment: str) -> bool:
        """Check if deployment has explicit configuration"""
        try:
            from chuk_llm.configuration import get_config

            config_manager = get_config()
            provider_config = config_manager.get_provider("azure_openai")

            # Check if deployment is in the models list
            if deployment in provider_config.models:
                return True

            # Check if any model capability pattern matches this deployment
            for capability in provider_config.model_capabilities:
                if hasattr(capability, "matches") and capability.matches(deployment):
                    return True
                elif hasattr(capability, "pattern"):
                    import re

                    if re.match(capability.pattern, deployment):
                        return True

            return False

        except Exception:
            return False

    def validate_model(self, model_name: str) -> bool:
        """
        Override model validation for Azure.

        KEY FIX: Azure deployments can have ANY name, so we ALWAYS return True
        for Azure OpenAI. This allows "scribeflowgpt4o" and any other custom
        deployment name to work! The actual validation happens when we try
        to use the deployment with the Azure API.
        """
        # For Azure, ANY deployment name is potentially valid
        # The actual validation happens when we try to use it
        log.debug(
            f"[azure_openai] Deployment '{model_name}' - allowing custom deployment (Azure supports any name)"
        )
        return True  # ALWAYS return True for Azure deployments

    def supports_feature(self, feature_name: str) -> bool:
        """
        Enhanced feature support with smart defaults for unknown Azure deployments.

        This is the KEY METHOD that enables discovered deployments to work!
        """
        try:
            # Get smart defaults first for this deployment
            smart_features = self._get_smart_default_features(self.azure_deployment)

            # First try the configuration system
            config_supports = super().supports_feature(feature_name)

            # If configuration gives True, trust it
            if config_supports is True:
                return True

            # If config says False but smart defaults say True, use smart defaults
            # This handles cases where config is missing/wrong for new deployments
            if config_supports is False and feature_name in smart_features:
                log.info(
                    f"[azure_openai] Config says no, but smart defaults say yes for '{self.azure_deployment}' - "
                    f"using smart default: supports {feature_name}"
                )
                return True

            # If config gave a definitive False and smart defaults also say False, trust it
            if config_supports is False:
                return False

            # Configuration returned None (unknown deployment) - use smart defaults
            # This is where we handle "scribeflowgpt4o" and other custom deployments!
            supports_smart = feature_name in smart_features

            if supports_smart:
                log.info(
                    f"[azure_openai] No config for deployment '{self.azure_deployment}' - "
                    f"using smart default: supports {feature_name}"
                )
            else:
                log.debug(
                    f"[azure_openai] No config for deployment '{self.azure_deployment}' - "
                    f"smart default: doesn't support {feature_name}"
                )

            return supports_smart

        except Exception as e:
            log.warning(f"Feature support check failed for {feature_name}: {e}")

            # For Azure, be optimistic about unknown features for chat models
            deployment_lower = self.azure_deployment.lower()
            if any(
                pattern in deployment_lower
                for pattern in ["gpt", "chat", "turbo", "o1", "o3"]
            ):
                log.info(
                    f"[azure_openai] Error checking config - assuming deployment '{self.azure_deployment}' "
                    f"supports {feature_name} (optimistic fallback)"
                )
                return True

            return False

    def get_model_info(self) -> dict[str, Any]:
        """
        Enhanced model info with smart defaults for discovered deployments.
        """
        # Get base info from configuration
        info = super().get_model_info()

        # Check if using smart defaults
        using_smart_defaults = not self._has_explicit_deployment_config(
            self.azure_deployment
        )

        # Get tool compatibility info
        tool_compatibility = self.get_tool_compatibility_info()

        # Add smart defaults info if applicable
        if using_smart_defaults:
            smart_features = self._get_smart_default_features(self.azure_deployment)
            smart_params = self._get_smart_default_parameters(self.azure_deployment)

            # Override with smart defaults
            info.update(
                {
                    "using_smart_defaults": True,
                    "smart_default_features": list(smart_features),
                    "smart_default_parameters": smart_params,
                    "discovery_note": f"Deployment '{self.azure_deployment}' not in static config - using smart defaults",
                    "features": list(smart_features),
                    "max_context_length": smart_params.get(
                        "max_context_length", 128000
                    ),
                    "max_output_tokens": smart_params.get("max_output_tokens", 8192),
                }
            )

        # Add Azure-specific metadata
        if not info.get("error"):
            info.update(
                {
                    "azure_specific": {
                        "endpoint": self.azure_endpoint,
                        "deployment": self.azure_deployment,
                        "api_version": self.api_version,
                        "authentication_type": self._get_auth_type(),
                        "deployment_to_model_mapping": True,
                        "supports_custom_deployments": True,  # Key feature!
                    },
                    "openai_compatible": True,
                    **tool_compatibility,
                    "parameter_mapping": {
                        "temperature": "temperature",
                        "max_tokens": "max_tokens",
                        "top_p": "top_p",
                        "frequency_penalty": "frequency_penalty",
                        "presence_penalty": "presence_penalty",
                        "stop": "stop",
                        "stream": "stream",
                        "tools": "tools",
                        "tool_choice": "tool_choice",
                    },
                    "azure_parameters": [
                        "azure_endpoint",
                        "api_version",
                        "azure_deployment",
                        "azure_ad_token",
                        "azure_ad_token_provider",
                    ],
                }
            )

        return info

    def _get_auth_type(self) -> str:
        """Determine the authentication type being used"""
        if (
            hasattr(self.async_client, "_azure_ad_token_provider")
            and self.async_client._azure_ad_token_provider
        ):
            return "azure_ad_token_provider"
        elif (
            hasattr(self.async_client, "_azure_ad_token")
            and self.async_client._azure_ad_token
        ):
            return "azure_ad_token"
        else:
            return "api_key"

    def _validate_request_with_config(
        self,
        messages: list[Message],
        tools: list[Any] | None = None,
        stream: bool = False,
        **kwargs,
    ) -> tuple[list[Message], list[Any] | None, bool, dict[str, Any]]:
        """
        Enhanced validation that uses smart defaults for unknown deployments.
        FIXED: Actually disable streaming when not supported.
        """
        # Import here to avoid circular imports
        from chuk_llm.core.models import Message as MessageModel
        from chuk_llm.core.models import Tool as ToolModel

        # Convert to Pydantic if needed
        def _convert_content_item(item):
            """Convert content item (for multimodal messages) to dict if needed."""
            if isinstance(item, dict):
                return item
            # Handle objects with attributes
            return dict(vars(item).items())

        def _convert_message(msg):
            if isinstance(msg, MessageModel):
                return msg
            elif isinstance(msg, dict):
                return MessageModel.model_validate(msg)
            else:
                # Handle objects with attributes (convert to dict first)
                content = getattr(msg, "content", None)
                # Convert content list items if needed
                if isinstance(content, list):
                    content = [_convert_content_item(item) for item in content]

                msg_dict = {
                    "role": getattr(msg, "role", None),
                    "content": content,
                }
                if hasattr(msg, "tool_calls") and msg.tool_calls:
                    msg_dict["tool_calls"] = msg.tool_calls
                if hasattr(msg, "tool_call_id") and msg.tool_call_id:
                    msg_dict["tool_call_id"] = msg.tool_call_id
                if hasattr(msg, "name") and msg.name:
                    msg_dict["name"] = msg.name
                return MessageModel.model_validate(msg_dict)

        validated_messages = [_convert_message(msg) for msg in messages]

        def _convert_tool(tool):
            if isinstance(tool, ToolModel):
                return tool
            elif isinstance(tool, dict):
                return ToolModel.model_validate(tool)
            else:
                # Handle objects with attributes
                return ToolModel.model_validate(
                    {"type": getattr(tool, "type", "function"), **vars(tool)}
                )

        validated_tools = [_convert_tool(tool) for tool in tools] if tools else None
        validated_stream = stream
        validated_kwargs = kwargs.copy()

        # Permissive approach: Don't block streaming
        # Let Azure OpenAI API handle unsupported cases - deployments can be added dynamically
        # and we shouldn't prevent attempts based on capability checks

        # Permissive approach: Pass all parameters to API (tools, vision, audio, JSON mode, etc.)
        # Let Azure OpenAI API handle unsupported cases - deployments can be added dynamically
        # Log when using smart defaults for unknown deployments
        if tools and not self._has_explicit_deployment_config(self.azure_deployment):
            log.info(
                f"Using smart default: assuming deployment '{self.azure_deployment}' supports tools"
            )

        # Get smart parameters for this deployment
        smart_params = self._get_smart_default_parameters(self.azure_deployment)

        # Apply parameter mapping (e.g., max_tokens -> max_completion_tokens for GPT-5/o1+ models)
        # This is applied REGARDLESS of whether there's explicit config, because it's based on model architecture
        parameter_mapping = smart_params.get("parameter_mapping", {})
        if parameter_mapping:
            for old_param, new_param in parameter_mapping.items():
                if old_param in validated_kwargs and new_param not in validated_kwargs:
                    original_value = validated_kwargs[old_param]
                    validated_kwargs[new_param] = original_value
                    validated_kwargs.pop(old_param)

                    # For reasoning models, ensure minimum max_completion_tokens
                    # Reasoning models need tokens for BOTH reasoning AND output
                    if new_param == "max_completion_tokens" and original_value < 1000:
                        # If user specified a small value like 150, bump it to at least 2000
                        # to allow room for reasoning tokens + output
                        validated_kwargs[new_param] = max(2000, original_value * 10)
                        log.debug(
                            f"[azure_openai] Increased max_completion_tokens from {original_value} to {validated_kwargs[new_param]} "
                            f"for reasoning model '{self.azure_deployment}' (needs tokens for reasoning + output)"
                        )
                    else:
                        log.debug(
                            f"[azure_openai] Mapped parameter '{old_param}' -> '{new_param}' for '{self.azure_deployment}'"
                        )

        # Remove unsupported parameters
        # This is applied REGARDLESS of whether there's explicit config, because it's based on model architecture
        unsupported_params = smart_params.get("unsupported_params", [])
        if unsupported_params:
            for param in unsupported_params:
                if param in validated_kwargs:
                    removed_value = validated_kwargs.pop(param)
                    log.debug(
                        f"[azure_openai] Removed unsupported parameter '{param}' (value: {removed_value}) for '{self.azure_deployment}'"
                    )

        # Apply smart parameter defaults only if no explicit config
        if not self._has_explicit_deployment_config(self.azure_deployment):
            # Apply smart defaults for max_tokens if not set
            if (
                "max_tokens" not in validated_kwargs
                and "max_completion_tokens" not in validated_kwargs
            ):
                if smart_params.get("requires_max_completion_tokens"):
                    validated_kwargs["max_completion_tokens"] = smart_params.get(
                        "max_output_tokens", 8192
                    )
                else:
                    validated_kwargs["max_tokens"] = smart_params.get(
                        "max_output_tokens", 8192
                    )

            log.debug(
                f"[azure_openai] Applied smart parameter defaults for '{self.azure_deployment}'"
            )

        # Validate parameters using configuration or smart defaults
        validated_kwargs = self.validate_parameters(**validated_kwargs)

        return validated_messages, validated_tools, validated_stream, validated_kwargs

    def _prepare_azure_request_params(self, **kwargs) -> dict[str, Any]:
        """Prepare request parameters for Azure OpenAI API"""
        # Use deployment name instead of model for Azure
        params = kwargs.copy()

        # Azure-specific parameter handling
        if "deployment_name" in params:
            params["model"] = params.pop("deployment_name")

        # Don't override if model is already set correctly
        if "model" not in params:
            params["model"] = self.azure_deployment

        return params

    # ------------------------------------------------------------------ #
    # Enhanced public API using configuration and smart defaults         #
    # ------------------------------------------------------------------ #
    def create_completion(
        self,
        messages: list[Message],
        tools: list[Tool] | None = None,
        *,
        stream: bool = False,
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any]] | Any:
        """
        Configuration-aware completion that validates capabilities before processing.
        Uses universal tool name compatibility with bidirectional mapping.
        Supports discovered deployments with smart defaults!
        """
        # Validate request against configuration (with smart defaults for unknown deployments)
        # Note: _validate_request_with_config handles Pydantic conversion internally
        validated_messages, validated_tools, validated_stream, validated_kwargs = (
            self._validate_request_with_config(messages, tools, stream, **kwargs)
        )

        # Apply universal tool name sanitization (stores mapping for restoration)
        # Keep tools as Pydantic objects throughout
        name_mapping = {}
        if validated_tools:
            validated_tools = self._sanitize_tool_names(validated_tools)
            name_mapping = self._current_name_mapping
            log.debug(
                f"Tool sanitization: {len(name_mapping)} tools processed for Azure OpenAI compatibility"
            )

        # Use configuration-aware parameter adjustment
        validated_kwargs = self._adjust_parameters_for_provider(validated_kwargs)

        # Convert Pydantic models to dicts for OpenAI SDK (only at final step)
        messages_dicts: list[dict[str, Any]] = [
            msg.model_dump() if hasattr(msg, "model_dump") else msg
            for msg in validated_messages
        ]
        # Use the helper method to convert tools to dicts at the final step
        tools_dicts: list[dict[str, Any]] | None = self._tools_to_dicts(validated_tools)

        if validated_stream:
            return self._stream_completion_async(
                messages_dicts, tools_dicts, name_mapping, **validated_kwargs
            )
        else:
            return self._regular_completion(
                messages_dicts, tools_dicts, name_mapping, **validated_kwargs
            )

    async def _stream_completion_async(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        name_mapping: dict[str, str] = None,
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any]]:
        """
        FIXED: Azure OpenAI streaming with proper JSON completion testing.

        Key fixes:
        - Only yield tool calls when JSON arguments are complete and parseable
        - Removed complex signature tracking system
        - Added completion status tracking
        - Prevents JSON parsing errors in downstream code
        """
        max_retries = 1

        for attempt in range(max_retries + 1):
            try:
                log.debug(
                    f"[azure_openai] Starting streaming (attempt {attempt + 1}): "
                    f"deployment={self.azure_deployment}, messages={len(messages)}, tools={len(tools) if tools else 0}"
                )

                # Prepare request parameters
                request_params = kwargs.copy()
                request_params["model"] = self.azure_deployment
                request_params["messages"] = messages
                if tools:
                    request_params["tools"] = tools
                request_params["stream"] = True

                # For reasoning models (GPT-5, o1+), enable streaming of reasoning tokens
                # This allows seeing the thinking process in real-time
                if self.supports_feature("reasoning"):
                    # Use stream_options to get reasoning content in the stream
                    if "stream_options" not in request_params:
                        request_params["stream_options"] = {"include_usage": True}
                    log.debug(
                        f"[azure_openai] Enabled stream_options for reasoning model '{self.azure_deployment}'"
                    )

                response_stream = await self.async_client.chat.completions.create(
                    **request_params
                )

                chunk_count = 0
                total_content = ""

                # FIXED: Simple completion-based tracking instead of signature system
                accumulated_tool_calls = {}  # {index: {id, name, arguments, complete}}

                async for chunk in response_stream:
                    chunk_count += 1

                    content = ""
                    completed_tool_calls = []  # Only completed tool calls this chunk

                    try:
                        if (
                            hasattr(chunk, "choices")
                            and chunk.choices
                            and len(chunk.choices) > 0
                        ):
                            choice = chunk.choices[0]

                            if hasattr(choice, "delta") and choice.delta:
                                delta = choice.delta

                                # Handle content - this works fine
                                if (
                                    hasattr(delta, "content")
                                    and delta.content is not None
                                ):
                                    content = str(delta.content)
                                    total_content += content

                                # FIXED: Handle tool calls with proper completion testing
                                if hasattr(delta, "tool_calls") and delta.tool_calls:
                                    for tc in delta.tool_calls:
                                        try:
                                            tc_index = getattr(tc, "index", 0)

                                            # Initialize accumulator with completion tracking
                                            if tc_index not in accumulated_tool_calls:
                                                accumulated_tool_calls[tc_index] = {
                                                    "id": getattr(
                                                        tc,
                                                        "id",
                                                        f"call_{uuid.uuid4().hex[:8]}",
                                                    ),
                                                    "name": "",
                                                    "arguments": "",
                                                    "complete": False,  # ADDED: Track completion status
                                                }

                                            tool_call_data = accumulated_tool_calls[
                                                tc_index
                                            ]

                                            # Update data
                                            if hasattr(tc, "id") and tc.id:
                                                tool_call_data["id"] = tc.id

                                            if hasattr(tc, "function") and tc.function:
                                                if (
                                                    hasattr(tc.function, "name")
                                                    and tc.function.name
                                                ):
                                                    tool_call_data["name"] += (
                                                        tc.function.name
                                                    )

                                                if (
                                                    hasattr(tc.function, "arguments")
                                                    and tc.function.arguments
                                                ):
                                                    tool_call_data["arguments"] += (
                                                        tc.function.arguments
                                                    )

                                            # CRITICAL FIX: Only yield when JSON is complete and valid
                                            if (
                                                tool_call_data["name"]
                                                and tool_call_data["arguments"]
                                                and not tool_call_data["complete"]
                                            ):
                                                try:
                                                    # Handle Azure-specific argument formatting
                                                    args_str = tool_call_data[
                                                        "arguments"
                                                    ]
                                                    if args_str.startswith(  # type: ignore[union-attr]
                                                        '""'  # type: ignore[union-attr]
                                                    ) and args_str.endswith('""'):  # type: ignore[union-attr]
                                                        args_str = args_str[2:-2]  # type: ignore[index]

                                                    # Test if JSON is complete and valid
                                                    parsed_args = json.loads(args_str)  # type: ignore[arg-type]

                                                    # Mark as complete and add to current chunk
                                                    tool_call_data["complete"] = True

                                                    tool_call = {
                                                        "id": tool_call_data["id"],
                                                        "type": "function",
                                                        "function": {
                                                            "name": tool_call_data[
                                                                "name"
                                                            ],
                                                            "arguments": json.dumps(
                                                                parsed_args
                                                            ),
                                                        },
                                                    }

                                                    completed_tool_calls.append(
                                                        tool_call
                                                    )
                                                    log.debug(
                                                        f"Azure tool call {tc_index} completed: {tool_call_data['name']}"
                                                    )

                                                except json.JSONDecodeError:
                                                    # JSON incomplete - keep accumulating
                                                    log.debug(
                                                        f"Azure tool call {tc_index} JSON incomplete, continuing accumulation"
                                                    )
                                                    pass

                                        except Exception as e:
                                            log.debug(
                                                f"Error processing Azure streaming tool call chunk: {e}"
                                            )
                                            continue

                    except Exception as chunk_error:
                        log.warning(
                            f"Error processing Azure chunk {chunk_count}: {chunk_error}"
                        )
                        content = ""

                    # Prepare result
                    result = {
                        "response": content,
                        "tool_calls": (
                            completed_tool_calls if completed_tool_calls else None
                        ),
                    }

                    # Restore tool names using universal restoration
                    if name_mapping and completed_tool_calls:
                        result = self._restore_tool_names_in_response(
                            result, name_mapping
                        )

                    # Only yield if we have content or completed tool calls
                    if content or completed_tool_calls:
                        yield result

                log.debug(
                    f"[azure_openai] Streaming completed: {chunk_count} chunks, "
                    f"{len(total_content)} total characters, {len(accumulated_tool_calls)} tool calls"
                )

                # If we reach here, streaming was successful - exit the retry loop
                return

            except Exception as e:
                error_str = str(e).lower()

                # Handle deployment errors immediately - these are not retryable
                if "deployment" in error_str and "not found" in error_str:
                    log.error(
                        f"[azure_openai] Deployment error - deployment '{self.azure_deployment}' not found: {e}"
                    )
                    yield {
                        "response": f"Azure deployment error: Deployment '{self.azure_deployment}' not found",
                        "tool_calls": [],
                        "error": True,
                    }
                    return  # Don't retry deployment errors

                # Check if this is a retryable error
                is_retryable = any(
                    pattern in error_str
                    for pattern in [
                        "timeout",
                        "connection",
                        "network",
                        "temporary",
                        "rate limit",
                    ]
                )

                if attempt < max_retries and is_retryable:
                    wait_time = (attempt + 1) * 1.0
                    log.warning(
                        f"[azure_openai] Streaming attempt {attempt + 1} failed: {e}. Retrying in {wait_time}s..."
                    )
                    await asyncio.sleep(wait_time)
                    continue  # Retry the request
                else:
                    # Final failure - yield exactly one error chunk
                    log.error(
                        f"[azure_openai] Streaming failed after {attempt + 1} attempts: {e}"
                    )
                    yield {
                        "response": f"Error: {str(e)}",
                        "tool_calls": [],
                        "error": True,
                    }
                    return

    async def _regular_completion(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        name_mapping: dict[str, str] = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Enhanced non-streaming completion using Azure OpenAI configuration with tool name restoration."""
        try:
            log.debug(
                f"[azure_openai] Starting completion: "
                f"deployment={self.azure_deployment}, messages={len(messages)}, tools={len(tools) if tools else 0}"
            )

            # Prepare request parameters - ensure model is set to deployment name
            request_params = kwargs.copy()
            request_params["model"] = self.azure_deployment
            request_params["messages"] = messages
            if tools:
                request_params["tools"] = tools
            request_params["stream"] = False

            resp = await self.async_client.chat.completions.create(**request_params)

            # Enhanced response debugging
            if hasattr(resp, "choices") and resp.choices:
                choice = resp.choices[0]
                log.debug(f"[azure_openai] Response choice type: {type(choice)}")
                if hasattr(choice, "message"):
                    message = choice.message
                    log.debug(f"[azure_openai] Message type: {type(message)}")
                    content_preview = getattr(message, "content", "NO CONTENT")
                    if content_preview:
                        log.debug(
                            f"[azure_openai] Content preview: {str(content_preview)[:100]}..."
                        )
                    else:
                        log.debug("[azure_openai] No content in message")

            # Use enhanced normalization from OpenAIStyleMixin
            result = self._normalize_message(resp.choices[0].message)

            # Restore original tool names using universal restoration
            if name_mapping and result.get("tool_calls"):
                result = self._restore_tool_names_in_response(result, name_mapping)

            # Log result
            log.debug(
                f"[azure_openai] Completion result: "
                f"response={len(str(result.get('response', ''))) if result.get('response') else 0} chars, "
                f"tool_calls={len(result.get('tool_calls', []))}"
            )

            return result

        except Exception as e:
            error_str = str(e).lower()

            # Check for deployment not found errors
            if "deployment" in error_str and "not found" in error_str:
                log.error(
                    f"[azure_openai] Deployment '{self.azure_deployment}' not found. "
                    f"Please check your Azure OpenAI resource for available deployments."
                )
                return {
                    "response": f"Deployment error: '{self.azure_deployment}' not found in your Azure OpenAI resource",
                    "tool_calls": [],
                    "error": True,
                }

            # Check for tool naming errors
            if "function" in error_str and (
                "name" in error_str or "invalid" in error_str
            ):
                log.error(
                    f"[azure_openai] Tool naming error (this should not happen with universal compatibility): {e}"
                )
                return {
                    "response": f"Tool naming error: {str(e)}",
                    "tool_calls": [],
                    "error": True,
                }

            log.error(f"[azure_openai] Error in completion: {e}")
            return {"response": f"Error: {str(e)}", "tool_calls": [], "error": True}

    def _normalize_message(self, msg) -> dict[str, Any]:
        """
        Azure-specific message normalization with FIXED argument parsing.

        CRITICAL FIX: Azure OpenAI returns tool arguments as JSON strings, not dicts.
        This properly handles the string format and ensures arguments are always
        properly formatted JSON strings for downstream processing.
        """
        try:
            # Use the inherited OpenAI normalization method as base
            result = super()._normalize_message(msg)

            # AZURE FIX: Ensure tool arguments are properly formatted JSON strings
            if result.get("tool_calls"):
                fixed_tool_calls = []

                for tool_call in result["tool_calls"]:
                    if "function" in tool_call and "arguments" in tool_call["function"]:
                        args = tool_call["function"]["arguments"]

                        # Azure often returns arguments as strings, sometimes double-quoted
                        if isinstance(args, str):
                            try:
                                # Try to parse the JSON to validate it
                                parsed_args = json.loads(args)
                                # Re-serialize to ensure consistent format
                                tool_call["function"]["arguments"] = json.dumps(
                                    parsed_args
                                )
                            except json.JSONDecodeError:
                                # If parsing fails, try to handle nested quoting
                                try:
                                    # Handle cases like ""{\"key\":\"value\"}"" (double quotes)
                                    if args.startswith('""') and args.endswith('""'):
                                        inner_json = args[2:-2]  # Remove outer quotes
                                        parsed_args = json.loads(inner_json)
                                        tool_call["function"]["arguments"] = json.dumps(
                                            parsed_args
                                        )
                                    else:
                                        # Last resort: wrap in empty object if invalid
                                        log.warning(
                                            f"Invalid tool arguments from Azure: {args}"
                                        )
                                        tool_call["function"]["arguments"] = "{}"
                                except Exception:
                                    tool_call["function"]["arguments"] = "{}"
                        elif isinstance(args, dict):
                            # Already a dict, convert to JSON string
                            tool_call["function"]["arguments"] = json.dumps(args)
                        else:
                            # Other types, default to empty object
                            tool_call["function"]["arguments"] = "{}"

                    fixed_tool_calls.append(tool_call)

                result["tool_calls"] = fixed_tool_calls

            return result

        except AttributeError:
            # Fallback implementation if mixin method not available
            content = None
            tool_calls = []

            # Extract content
            if hasattr(msg, "content"):
                content = msg.content

            # Extract tool calls with FIXED argument handling
            if hasattr(msg, "tool_calls") and msg.tool_calls:
                for tc in msg.tool_calls:
                    try:
                        # Get function arguments - Azure specific handling
                        raw_args = getattr(tc.function, "arguments", "{}")

                        # AZURE FIX: Properly handle argument formats
                        if isinstance(raw_args, str):
                            try:
                                # Validate JSON and reformat
                                parsed_args = json.loads(raw_args)
                                formatted_args = json.dumps(parsed_args)
                            except json.JSONDecodeError:
                                log.warning(
                                    f"Invalid JSON in Azure tool call arguments: {raw_args}"
                                )
                                formatted_args = "{}"
                        elif isinstance(raw_args, dict):
                            formatted_args = json.dumps(raw_args)
                        else:
                            log.warning(
                                f"Unexpected Azure argument type: {type(raw_args)}"
                            )
                            formatted_args = "{}"

                        tool_calls.append(
                            {
                                "id": getattr(tc, "id", f"call_{uuid.uuid4().hex[:8]}"),
                                "type": "function",
                                "function": {
                                    "name": getattr(tc.function, "name", "unknown"),
                                    "arguments": formatted_args,  # Always a properly formatted JSON string
                                },
                            }
                        )
                    except Exception as e:
                        log.warning(f"Failed to process Azure tool call: {e}")
                        continue

            # Return standard format
            if tool_calls:
                return {
                    "response": content if content else None,
                    "tool_calls": tool_calls,
                }
            else:
                return {"response": content or "", "tool_calls": []}

    def _adjust_parameters_for_provider(self, kwargs: dict[str, Any]) -> dict[str, Any]:
        """
        Adjust parameters using configuration instead of hardcoded rules.
        Now with smart defaults for unknown deployments!
        """
        adjusted = kwargs.copy()

        try:
            # Get smart parameters for this deployment
            smart_params = self._get_smart_default_parameters(self.azure_deployment)

            # Apply parameter mapping (e.g., max_tokens -> max_completion_tokens)
            # This must be done BEFORE validation to ensure correct parameter names
            parameter_mapping = smart_params.get("parameter_mapping", {})
            if parameter_mapping:
                for old_param, new_param in parameter_mapping.items():
                    if old_param in adjusted and new_param not in adjusted:
                        original_value = adjusted[old_param]
                        adjusted[new_param] = original_value
                        adjusted.pop(old_param)

                        # For reasoning models, ensure minimum max_completion_tokens
                        if (
                            new_param == "max_completion_tokens"
                            and original_value < 1000
                        ):
                            adjusted[new_param] = max(2000, original_value * 10)
                            log.debug(
                                f"[azure_openai] Increased max_completion_tokens from {original_value} to {adjusted[new_param]} "
                                f"for reasoning model (needs tokens for reasoning + output)"
                            )
                        else:
                            log.debug(
                                f"[azure_openai] Mapped parameter '{old_param}' -> '{new_param}' in adjust_parameters"
                            )

            # Remove unsupported parameters
            unsupported_params = smart_params.get("unsupported_params", [])
            if unsupported_params:
                for param in unsupported_params:
                    if param in adjusted:
                        removed_value = adjusted.pop(param)
                        log.debug(
                            f"[azure_openai] Removed unsupported parameter '{param}' (value: {removed_value}) in adjust_parameters"
                        )

            # Use the configuration-aware parameter validation
            adjusted = self.validate_parameters(**adjusted)

            # If no explicit config, use smart defaults
            if not self._has_explicit_deployment_config(self.azure_deployment):
                # Apply smart defaults if not already set
                if (
                    "max_tokens" not in adjusted
                    and "max_completion_tokens" not in adjusted
                ):
                    max_output = smart_params.get("max_output_tokens", 4096)
                    if smart_params.get("requires_max_completion_tokens"):
                        adjusted["max_completion_tokens"] = max_output
                    else:
                        adjusted["max_tokens"] = max_output
                    log.debug(
                        f"[azure_openai] Applied smart default for max output tokens={max_output} for '{self.azure_deployment}'"
                    )

            # Additional Azure OpenAI-specific parameter handling
            model_caps = self._get_model_capabilities()
            if model_caps:
                # Adjust max_tokens based on config if not already handled
                if "max_tokens" in adjusted and model_caps.max_output_tokens:
                    if adjusted["max_tokens"] > model_caps.max_output_tokens:
                        log.debug(
                            f"Adjusting max_tokens from {adjusted['max_tokens']} to {model_caps.max_output_tokens} for azure_openai"
                        )
                        adjusted["max_tokens"] = model_caps.max_output_tokens

                # Also adjust max_completion_tokens for GPT-5 and reasoning models
                if "max_completion_tokens" in adjusted and model_caps.max_output_tokens:
                    if adjusted["max_completion_tokens"] > model_caps.max_output_tokens:
                        log.debug(
                            f"Adjusting max_completion_tokens from {adjusted['max_completion_tokens']} to {model_caps.max_output_tokens} for azure_openai"
                        )
                        adjusted["max_completion_tokens"] = model_caps.max_output_tokens

        except Exception as e:
            log.debug(f"Could not adjust parameters using config: {e}")
            # Fallback: check if we need max_completion_tokens instead
            smart_params = self._get_smart_default_parameters(self.azure_deployment)
            if smart_params.get("requires_max_completion_tokens"):
                if "max_completion_tokens" not in adjusted:
                    adjusted["max_completion_tokens"] = 4096
            else:
                if "max_tokens" not in adjusted:
                    adjusted["max_tokens"] = 4096

        return adjusted

    async def close(self):
        """Cleanup resources"""
        # Reset name mapping
        self._current_name_mapping = {}
        try:
            if hasattr(self.async_client, "close"):
                await self.async_client.close()
        except RuntimeError as e:
            # Event loop may already be closed, which is fine
            if "Event loop is closed" not in str(e):
                raise
        try:
            if hasattr(self.client, "close"):
                self.client.close()
        except RuntimeError as e:
            # Event loop may already be closed, which is fine
            if "Event loop is closed" not in str(e):
                raise

    def __repr__(self) -> str:
        return f"AzureOpenAILLMClient(deployment={self.azure_deployment}, model={self.model}, endpoint={self.azure_endpoint})"
