#
# Copyright 2026 The Dapr Authors
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import logging
from typing import Any, Dict, Optional, Union

from openai import AzureOpenAI, OpenAI
from pydantic import Field

from dapr_agents.llm.base import LLMClientBase
from dapr_agents.llm.openai.client.azure import AzureOpenAIClient
from dapr_agents.llm.openai.client.openai import OpenAIClient
from dapr_agents.types.llm import AzureOpenAIClientConfig, OpenAIClientConfig

logger = logging.getLogger(__name__)


class OpenAIClientBase(LLMClientBase):
    """
    Base class for managing OpenAI and Azure OpenAI clients.
    Handles client initialization, configuration, and shared logic.
    """

    api_key: Optional[str] = Field(
        default=None, description="API key for OpenAI or Azure OpenAI."
    )
    base_url: Optional[str] = Field(
        default=None, description="Base URL for OpenAI API (OpenAI-specific)."
    )
    azure_endpoint: Optional[str] = Field(
        default=None, description="Azure endpoint URL (Azure OpenAI-specific)."
    )
    azure_deployment: Optional[str] = Field(
        default=None, description="Azure deployment name (Azure OpenAI-specific)."
    )
    api_version: Optional[str] = Field(
        default=None, description="Azure API version (Azure OpenAI-specific)."
    )
    organization: Optional[str] = Field(
        default=None, description="Organization for OpenAI or Azure OpenAI."
    )
    project: Optional[str] = Field(
        default=None, description="Project for OpenAI or Azure OpenAI."
    )
    azure_ad_token: Optional[str] = Field(
        default=None, description="Azure AD token for authentication (Azure-specific)."
    )
    azure_client_id: Optional[str] = Field(
        default=None, description="Client ID for Managed Identity (Azure-specific)."
    )
    timeout: Union[int, float, Dict[str, Any]] = Field(
        default=1500, description="Timeout for requests in seconds."
    )

    def model_post_init(self, __context: Any) -> None:
        """
        Initializes private attributes after validation.
        """
        self._provider = "openai"

        # Set up the private config and client attributes
        self._config: Union[AzureOpenAIClientConfig, OpenAIClientConfig] = (
            self.get_config()
        )
        self._client: Union[AzureOpenAI, OpenAI] = self.get_client()
        return super().model_post_init(__context)

    def get_config(self) -> Union[OpenAIClientConfig, AzureOpenAIClientConfig]:
        """
        Returns the appropriate configuration for OpenAI or Azure OpenAI.
        """
        is_azure = self.azure_endpoint or self.azure_deployment

        if is_azure:
            return AzureOpenAIClientConfig(
                api_key=self.api_key,
                organization=self.organization,
                project=self.project,
                azure_ad_token=self.azure_ad_token,
                azure_endpoint=self.azure_endpoint,
                azure_deployment=self.azure_deployment,
                api_version=self.api_version,
            )
        else:
            return OpenAIClientConfig(
                api_key=self.api_key,
                base_url=self.base_url,
                organization=self.organization,
                project=self.project,
            )

    def get_client(self) -> Union[AzureOpenAI, OpenAI]:
        """
        Initialize and return the appropriate client (OpenAI or Azure OpenAI).
        """
        config = self.config
        timeout = self.timeout

        if isinstance(config, AzureOpenAIClientConfig):
            logger.info("Initializing Azure OpenAI client...")
            return AzureOpenAIClient(
                api_key=config.api_key,
                azure_ad_token=config.azure_ad_token,
                azure_endpoint=config.azure_endpoint,
                azure_deployment=config.azure_deployment,
                api_version=config.api_version,
                organization=config.organization,
                project=config.project,
                azure_client_id=self.azure_client_id,
                timeout=timeout,
            ).get_client()

        logger.info("Initializing OpenAI client...")
        return OpenAIClient(
            api_key=config.api_key,
            base_url=config.base_url,
            organization=config.organization,
            project=config.project,
            timeout=timeout,
        ).get_client()

    @property
    def config(self) -> Union[AzureOpenAIClientConfig, OpenAIClientConfig]:
        return self._config

    @property
    def client(self) -> Union[OpenAI, AzureOpenAI]:
        return self._client