"""OpenAI LLM provider implementation."""

from collections.abc import AsyncGenerator
from typing import Any

from loguru import logger
from openai import AsyncOpenAI, OpenAI
from openai.types.chat.chat_completion import ChatCompletion

from kguru.config import OpenAIConfig
from kguru.core.reasoning.llm.base import BaseLLM
from kguru.utils import async_retry

DEFAULT_MAX_TOKENS = 4096


class OpenAILLM(BaseLLM):
    """OpenAI LLM provider implementation."""

    reasoning_models = ["o1", "o1-mini", "o1-preview", "o3-mini"]

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

        Args:
            model: Model identifier (e.g., "gpt-4", "gpt-4-turbo")
            api_key: OpenAI API key
            api_base: Optional base URL for API calls
            **additional_kwargs: Additional OpenAI-specific arguments
        """
        super().__init__(
            model=model,
            api_key=api_key,
            api_base=api_base,
            **additional_kwargs,
        )

    @property
    def client(self) -> OpenAI:
        """Get synchronous OpenAI client."""
        if self._client is None:
            kwargs = {"api_key": self.api_key}
            if self.api_base:
                kwargs["base_url"] = self.api_base
            self._client = OpenAI(**kwargs)
        return self._client

    @property
    def aclient(self) -> AsyncOpenAI:
        """Get async OpenAI client."""
        if self._aclient is None:
            kwargs = {"api_key": self.api_key}
            if self.api_base:
                kwargs["base_url"] = self.api_base
            self._aclient = AsyncOpenAI(**kwargs)
        return self._aclient

    def validate_configuration(self) -> None:
        """Validate 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) -> "OpenAILLM":
        """Get default OpenAI LLM instance."""
        return cls()

    @classmethod
    def get_reasoning_llm(cls, model: str = "o1-mini") -> "OpenAILLM":
        """
        Get OpenAI reasoning model instance.

        Args:
            model: Reasoning model identifier

        Returns:
            OpenAI LLM instance configured for reasoning
        """
        logger.debug(f"Using {model} for reasoning")
        return cls(model=model)

    @classmethod
    def get_light_llm(cls, model: str = "gpt-4-turbo") -> "OpenAILLM":
        """
        Get lightweight OpenAI model instance.

        Args:
            model: Lightweight model identifier

        Returns:
            OpenAI LLM instance configured for lightweight tasks
        """
        logger.debug(f"Using {model} for lightweight tasks")
        return cls(model=model)

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

        Args:
            **kwargs: OpenAI chat completion arguments

        Returns:
            Raw ChatCompletion response
        """
        # Set default temperature to 0 if not provided (except for reasoning models)
        if "temperature" not in kwargs and self.model not in self.reasoning_models:
            kwargs["temperature"] = 0
        if "max_tokens" not in kwargs and self.model not in self.reasoning_models:
            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 OpenAI.

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

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

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

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

        Yields:
            Text chunks as they are generated
        """
        # Set default temperature to 0 if not provided (except for reasoning models)
        if "temperature" not in kwargs and self.model not in self.reasoning_models:
            kwargs["temperature"] = 0
        if "max_tokens" not in kwargs and self.model not in self.reasoning_models:
            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
