"""Azure OpenAI LLM provider implementation."""

from collections.abc import AsyncGenerator
from typing import Any

from loguru import logger
from openai import AsyncAzureOpenAI, AzureOpenAI
from openai.types.chat.chat_completion import ChatCompletion

from kguru.config import AzureOpenAIConfig
from kguru.core.reasoning.llm.base import BaseLLM
from kguru.utils import async_retry, extract_json

DEFAULT_MAX_TOKENS = 4096


class AzureOpenAILLM(BaseLLM):
    """Azure OpenAI LLM provider implementation."""

    def __init__(
        self,
        model: str | None = None,
        api_key: str | None = None,
        azure_endpoint: str | None = None,
        api_version: str | None = None,
        deployment: str | None = None,
        **additional_kwargs: Any,
    ) -> None:
        """
        Initialize Azure OpenAI LLM provider.

        Args:
            model: Model identifier
            api_key: Azure OpenAI API key
            azure_endpoint: Azure OpenAI endpoint URL
            api_version: Azure OpenAI API version
            deployment: Azure deployment name
            **additional_kwargs: Additional Azure-specific arguments
        """
        # Use config defaults if not provided
        final_model = model if model is not None else AzureOpenAIConfig.model
        final_api_key = api_key if api_key is not None else AzureOpenAIConfig.api_key
        final_azure_endpoint = (
            azure_endpoint if azure_endpoint is not None else AzureOpenAIConfig.azure_endpoint
        )
        final_api_version = (
            api_version if api_version is not None else AzureOpenAIConfig.api_version
        )
        final_deployment = deployment if deployment is not None else AzureOpenAIConfig.deployment

        super().__init__(
            model=final_model,
            api_key=final_api_key,
            api_base=final_azure_endpoint,
            **additional_kwargs,
        )
        # Azure-specific attributes
        self.deployment = final_deployment
        self.api_version = final_api_version
        self.azure_endpoint = final_azure_endpoint

    @property
    def client(self) -> AzureOpenAI:
        """Get synchronous Azure OpenAI client."""
        if self._client is None:
            self._client = AzureOpenAI(
                api_key=self.api_key,
                azure_endpoint=self.azure_endpoint,
                azure_deployment=self.deployment,
                api_version=self.api_version,
            )
        return self._client

    @property
    def aclient(self) -> AsyncAzureOpenAI:
        """Get async Azure OpenAI client."""
        if self._aclient is None:
            self._aclient = AsyncAzureOpenAI(
                api_key=self.api_key,
                azure_endpoint=self.azure_endpoint,
                azure_deployment=self.deployment,
                api_version=self.api_version,
            )
        return self._aclient

    def validate_configuration(self) -> None:
        """Validate Azure OpenAI configuration by making a test request."""
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": "Hello, world!"}],
            )
            if response.choices[0].message.content is None:
                raise ValueError("Failed to validate configuration")
            logger.debug(f"Validation response: {response.choices[0].message.content}")
            logger.debug("Configuration validated successfully")
        except Exception as e:
            logger.error(f"Configuration validation failed: {e}")
            raise

    @classmethod
    def get_default(cls) -> "AzureOpenAILLM":
        """Get default Azure OpenAI LLM instance."""
        return cls()

    @async_retry(max_retries=3, delay=2)
    async def _raw_completion(self, **kwargs: Any) -> ChatCompletion:
        """
        Make raw Azure OpenAI API call.

        Args:
            **kwargs: Azure OpenAI chat completion arguments

        Returns:
            Raw ChatCompletion response
        """
        # Set default temperature to 0 if not provided
        if "temperature" not in kwargs:
            kwargs["temperature"] = 0
        if "max_tokens" not in kwargs:
            kwargs["max_tokens"] = DEFAULT_MAX_TOKENS

        return await self.aclient.chat.completions.create(model=self.model, **kwargs)

    async def reason(self, **kwargs: Any) -> str:
        """
        Generate text response using Azure OpenAI.

        Args:
            **kwargs: Azure OpenAI chat completion arguments (messages, temperature, etc.)

        Returns:
            Generated text response
        """
        result = await self._raw_completion(**kwargs)
        content = result.choices[0].message.content
        return content if content is not None else ""

    async def reason_structure(self, **kwargs: Any) -> dict[str, Any]:
        """
        Generate structured JSON response using Azure OpenAI.

        Args:
            **kwargs: Azure OpenAI chat completion arguments (messages, temperature, etc.)

        Returns:
            Parsed JSON response as dictionary
        """
        result = await self._raw_completion(**kwargs)
        content = result.choices[0].message.content
        text: str = content if content is not None else ""
        return extract_json(text)

    async def stream(self, **kwargs: Any) -> AsyncGenerator[str, None]:
        """
        Stream text response chunks from Azure OpenAI.

        Args:
            **kwargs: Azure OpenAI chat completion arguments (messages, temperature, etc.)

        Yields:
            Text chunks as they are generated
        """
        # Set default temperature to 0 if not provided
        if "temperature" not in kwargs:
            kwargs["temperature"] = 0
        if "max_tokens" not in kwargs:
            kwargs["max_tokens"] = DEFAULT_MAX_TOKENS

        # Enable streaming
        kwargs["stream"] = True

        # Create streaming response
        stream = await self.aclient.chat.completions.create(model=self.model, **kwargs)

        # Yield content chunks
        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content is not None:
                yield chunk.choices[0].delta.content
