"""
Abstract base class for LLM chat completion clients.

This module defines the interface that all LLM provider integrations must implement.
Provides a unified way to interact with different language models (OpenAI, Anthropic, etc.)
while abstracting away provider-specific API differences.

Design Pattern:
    Every provider implementation follows the Three-Step Conversion Pattern:
    1. Convert Input: Internal Message types → Provider-specific format
    2. Make API Call: Send to provider with authentication
    3. Convert Output: Provider response → Unified ChatCompletionResult
"""

import json
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator
from typing import Any, Dict, List, Optional, TypeVar

from pydantic import BaseModel

from agentbyte.component import ComponentLoader
from agentbyte.messages import AssistantMessage, Message, MultiModalMessage, ToolMessage

from .types import (
    ChatCompletionChunk,
    ChatCompletionResult,
    ModelClientError,
)

ConfigT = TypeVar("ConfigT", bound="BaseChatCompletionClientConfig")


class BaseChatCompletionClientConfig(BaseModel):
    """
    Base configuration class for LLM client serialization.

    This class serves as the base for provider-specific configurations,
    enabling round-trip serialization and deserialization of clients
    for enterprise scenarios like:
    - Configuration file management
    - Cross-environment deployment
    - Integration with orchestration frameworks
    - Configuration validation via Pydantic

    Provider implementations should:
    1. Subclass this with provider-specific fields
    2. Convert client state to config via _to_config() method
    3. Reconstruct client from config via _from_config() classmethod

    Example:
        ```python
        from agentbyte.llm import RetryPolicy

        class OpenAIChatCompletionClientConfig(BaseChatCompletionClientConfig):
            model: str
            config: Dict[str, Any] = {}
            retry_policy: RetryPolicy = RetryPolicy()

        # Serialize: client → config → JSON
        config = client._to_config()
        config_json = config.model_dump_json()

        # Deserialize: JSON → config → client
        loaded_config = OpenAIChatCompletionClientConfig.model_validate_json(config_json)
        new_client = OpenAIChatCompletionClient._from_config(loaded_config)
        ```
    """

    model: str = "gpt-4.1-mini"
    config: Dict[str, Any] = {}


class BaseChatCompletionClient(ABC, ComponentLoader):
    """
    Abstract base class for LLM chat completion client.

    Defines the interface that all LLM provider integrations must implement.
    This enables seamless switching between providers (OpenAI, Anthropic, Azure, etc.)
    without changing agent code.

    Design Pattern - Dependency Injection:
        The actual HTTP client (e.g., AsyncOpenAI, AsyncAnthropic) is injected at
        initialization. This allows organizations to:
        - Use different authentication methods (API keys, OAuth, service principals, etc.)
        - Support multiple environments with different credentials
        - Configure clients with proxies, custom endpoints, timeouts, etc.
        - Enable easy testing with mock clients
        - Manage client lifecycle independently

    Key Design Decisions:
        - Async-first: LLM calls are I/O-bound and should not block other tasks
        - Streaming support: Enable real-time updates and reduced perceived latency
        - Tool/function calling: Built-in support for tool invocation patterns
        - Unified error handling: All provider errors convert to ModelClientError
        - Dependency injection: Client passed in at initialization

    Example - OpenAI with API Key:
        >>> from openai import AsyncOpenAI
        >>> from agentbyte.llm import OpenAIChatCompletionClient, UserMessage
        >>> 
        >>> # Create configured client with API key
        >>> openai_client = AsyncOpenAI(api_key="sk-...")
        >>> 
        >>> # Inject into agentbyte
        >>> llm = OpenAIChatCompletionClient(
        ...     model="gpt-4.1-mini",
        ...     client=openai_client
        ... )
        >>> 
        >>> result = await llm.create(
        ...     messages=[UserMessage(content="Hello")]
        ... )

    Example - OpenAI with Custom Configuration:
        >>> from openai import AsyncOpenAI
        >>> from agentbyte.llm import OpenAIChatCompletionClient
        >>> 
        >>> # Custom client with proxy, timeout, etc.
        >>> openai_client = AsyncOpenAI(
        ...     api_key="sk-...",
        ...     base_url="https://proxy.company.com/openai",
        ...     timeout=60.0,
        ...     organization="org-123"
        ... )
        >>> llm = OpenAIChatCompletionClient(
        ...     model="gpt-4.1-mini",
        ...     client=openai_client
        ... )

    Example - OAuth or Custom Authentication:
        >>> from my_company import OAuthClient
        >>> 
        >>> # Use any authentication method
        >>> oauth_client = OAuthClient(
        ...     token_endpoint="https://auth.company.com/token",
        ...     client_id="...",
        ...     client_secret="..."
        ... )
        >>> 
        >>> llm = CustomChatCompletionClient(
        ...     model="gpt-4",
        ...     client=oauth_client
        ... )

    Example - Sync/Async Support:
        >>> # For organizations with sync APIs
        >>> # Create your own wrapper that adapts to the interface
        >>> class SyncToAsyncAdapter:
        ...     def __init__(self, sync_client):
        ...         self.sync_client = sync_client
        ...     
        ...     async def chat_completions_create(self, **kwargs):
        ...         # Run sync call in executor
        ...         loop = asyncio.get_event_loop()
        ...         return await loop.run_in_executor(
        ...             None,
        ...             lambda: self.sync_client.chat.completions.create(**kwargs)
        ...         )
        >>>
        >>> sync_client = OpenAI(api_key="sk-...")
        >>> adapter = SyncToAsyncAdapter(sync_client)
        >>> llm = OpenAIChatCompletionClient(
        ...     model="gpt-4.1-mini",
        ...     client=adapter
        ... )
    """

    def __init__(
        self,
        model: str,
        client: Any,
        config: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ):
        """
        Initialize the chat completion client with dependency injection.

        Args:
            model: The model identifier (e.g., "gpt-4.1-mini", "claude-3-sonnet").
                   This identifies which model to use for completions.

            client: Dependency-injected HTTP client (e.g., AsyncOpenAI, AsyncAnthropic).
                   This allows organizations to:
                   - Use any authentication method (API keys, OAuth, etc.)
                   - Configure proxies, timeouts, custom endpoints
                   - Support multiple environments with different credentials
                   - Enable testing with mock clients
                   
                   Provider implementations will use this client to make API calls.

            config: Optional dictionary of default configuration options.
                   These can be overridden per-request via **kwargs in create() calls.
                   Examples:
                   - temperature: Sampling temperature (0-2)
                   - max_tokens: Maximum tokens to generate
                   - top_p: Nucleus sampling parameter
                   - frequency_penalty: Penalty for repeated tokens
                   - presence_penalty: Penalty for new tokens
                   - provider-specific options

            **kwargs: Additional configuration options that override config dict.
                     Same options as config parameter.

        Example:
            ```python
            from openai import AsyncOpenAI
            
            # Create configured client with desired auth method
            my_client = AsyncOpenAI(
                api_key="sk-...",
                base_url="https://...",
                organization="org-123"
            )
            
            # Inject into agentbyte
            llm = OpenAIChatCompletionClient(
                model="gpt-4.1-mini",
                client=my_client,
                config={"temperature": 0.7, "max_tokens": 1000}
            )
            ```
        """
        self.model = model
        self.client = client  # Dependency injected
        self.config = config or {}
        # Additional kwargs override config
        self.config.update(kwargs)

    @abstractmethod
    async def create(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        """
        Make a single LLM API call.

        This is the core method for non-streaming interaction with an LLM.
        Follows the Three-Step Conversion Pattern:
        1. Convert input messages to provider format
        2. Make the API call using the injected client with tools if provided
        3. Parse response and convert back to unified format

        Implementation Notes:
            - Use self.client (the injected HTTP client) to make API calls
            - Merge config from self.config and **kwargs (kwargs override config)
            - Convert internal Message types to provider format via _convert_messages_to_api_format()
            - Convert provider response back to ChatCompletionResult
            - Catch provider-specific exceptions and convert to ModelClientError hierarchy

        Args:
            messages: List of messages in conversation history.
                First message is typically SystemMessage with instructions.
                Subsequent messages alternate between UserMessage and AssistantMessage,
                with ToolMessage interspersed for tool execution results.

            tools: Optional list of tool/function schemas available to the LLM.
                Each tool dict should have: name, description, parameters.
                Provider implementations convert to provider-specific format.
                Example:
                    ```
                    {
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "description": "Get weather for a location",
                            "parameters": {
                                "type": "object",
                                "properties": {
                                    "location": {"type": "string"}
                                },
                                "required": ["location"]
                            }
                        }
                    }
                    ```

            output_format: Optional Pydantic model class used for structured output.
                Provider implementations may pass JSON schema to the model and
                parse responses into structured objects.

            **kwargs: Provider-specific parameters that override self.config.
                - temperature: Sampling temperature (0-2, higher = more random)
                - max_tokens: Maximum tokens to generate
                - top_p: Nucleus sampling parameter
                - frequency_penalty: Penalty for repeated tokens
                - presence_penalty: Penalty for new tokens
                - etc. (provider-specific)

        Returns:
            ChatCompletionResult: Structured response containing:
                - message: AssistantMessage with content and optional tool_calls
                - usage: Token consumption (input, output, total)
                - model: Actual model used (may differ from requested)
                - metadata: Provider-specific data (finish_reason, etc.)
                - structured_output: Parsed output if output_format was provided

        Raises:
            ModelClientError: Base exception for all LLM client errors.
                May be subclassed to:
                - RateLimitError: API rate limits exceeded
                - AuthenticationError: Invalid API key or credentials
                - InvalidRequestError: Request format is invalid
                - APIError: Generic API error

        Example:
            ```python
            from openai import AsyncOpenAI
            from agentbyte.llm import (
                OpenAIChatCompletionClient,
                SystemMessage,
                UserMessage
            )

            # Create and configure the underlying client
            openai_client = AsyncOpenAI(api_key="sk-...", organization="org-123")

            # Inject into agentbyte
            llm = OpenAIChatCompletionClient(
                model="gpt-4.1-mini",
                client=openai_client,
                config={"temperature": 0.7}
            )

            # Use it
            result = await llm.create(
                messages=[
                    SystemMessage(content="You are helpful."),
                    UserMessage(content="What is 2+2?")
                ],
                temperature=0.5  # Override default config
            )

            print(result.message.content)  # "2 + 2 equals 4"
            print(f"Tokens: {result.usage.total_tokens}")
            ```
        """
        pass

    @abstractmethod
    async def create_stream(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        """
        Make a streaming LLM API call.

        Yields response chunks as they arrive from the API, enabling real-time updates
        and reduced time-to-first-token perception. Particularly useful for:
        - Long-running operations where immediate feedback improves UX
        - User-facing chat applications
        - Monitoring token generation in real-time

        Implementation Notes:
            - Use self.client (the injected HTTP client) to make streaming API calls
            - Merge config from self.config and **kwargs (kwargs override config)
            - Enable streaming mode on the underlying client call
            - Yield chunks as they arrive from the provider
            - Catch provider-specific exceptions and convert to ModelClientError hierarchy

        Args:
            messages: List of messages (same format as create())
            tools: Optional tool schemas (same format as create())
            output_format: Optional Pydantic model class (kept for API consistency;
                may be ignored by streaming providers)
            **kwargs: Provider-specific parameters that override self.config
                     (same as create())

        Yields:
            ChatCompletionChunk: Streaming response chunks containing:
                - content: Partial text content from this chunk (may be empty)
                - model: Model name
                - metadata: Provider-specific data (finish_reason for final chunk)

        Raises:
            ModelClientError: Same exception hierarchy as create()

        Example:
            ```python
            from openai import AsyncOpenAI
            from agentbyte.llm import OpenAIChatCompletionClient, UserMessage

            # Create and configure the underlying client
            openai_client = AsyncOpenAI(api_key="sk-...")

            # Inject into agentbyte
            llm = OpenAIChatCompletionClient(
                model="gpt-4.1-mini",
                client=openai_client
            )

            # Stream real-time updates
            async for chunk in llm.create_stream(
                messages=[UserMessage(content="Write a haiku")]
            ):
                print(chunk.content, end="", flush=True)
            print()  # Newline after done
            ```

        Note:
            - Chunks arrive incrementally; accumulate .content to rebuild message
            - The final chunk may have empty content but includes finish_reason
            - Use chunk count to detect API issues (too few/many chunks)
            - Some providers may emit empty chunks between content; handle gracefully
        """
        pass

    @abstractmethod
    def _to_config(self) -> BaseChatCompletionClientConfig:
        """
        Convert client instance to configuration for serialization.

        This method extracts the client's current state into a configuration object
        that can be serialized to JSON or stored in a config file.

        Returns:
            Provider-specific config subclass instance with all necessary state

        Example:
            ```python
            client = OpenAIChatCompletionClient(
                model="gpt-4.1-mini",
                client=openai_client,
                config={"temperature": 0.7}
            )
            config = client._to_config()
            # config now has: model, config, retry_policy (nested RetryPolicy), etc.
            ```
        """
        pass

    @classmethod
    @abstractmethod
    def _from_config(cls, config: BaseChatCompletionClientConfig) -> "BaseChatCompletionClient":
        """
        Reconstruct client from configuration.

        This classmethod should validate the config and recreate the client
        with the saved state.

        Args:
            config: Provider-specific config instance

        Returns:
            Reconstructed client ready for use

        Raises:
            ValueError: If config is invalid for this provider
            ImportError: If required dependencies are missing

        Example:
            ```python
            config = OpenAIChatCompletionClientConfig(
                model="gpt-4.1-mini",
                config={"temperature": 0.7}
            )
            client = OpenAIChatCompletionClient._from_config(config)
            ```
        """
        pass

    def _convert_messages_to_api_format(
        self, messages: List[Message]
    ) -> List[Dict[str, Any]]:
        """
        Convert Agentbyte message types to provider-specific API format.

        Helper method that provider implementations should override as needed.
        Handles the common case but subclasses may need custom logic for:
        - Vision/image content
        - Special role mappings
        - Tool call formatting differences
        - Prompt caching or other optimizations

        Default Implementation:
            Converts each Agentbyte Message to a dict with:
            - role: "system", "user", "assistant", or "tool"
            - content: Message text content
            - tool_calls: [list of tool calls] (only for AssistantMessage)
            - tool_call_id: str (only for ToolMessage, links to ToolCall.id)

        Args:
            messages: List of Agentbyte Message objects

        Returns:
            List of dicts in provider API format ready for API call

        Example Output:
            ```python
            [
                {"role": "system", "content": "You are helpful."},
                {"role": "user", "content": "What is 2+2?"},
                {
                    "role": "assistant",
                    "content": "Let me calculate that.",
                    "tool_calls": [
                        {
                            "id": "call_123",
                            "type": "function",
                            "function": {
                                "name": "calculator",
                                "arguments": "{\"expression\": \"2+2\"}"
                            }
                        }
                    ]
                },
                {
                    "role": "tool",
                    "content": "4",
                    "tool_call_id": "call_123"
                }
            ]
            ```
        """
        api_messages = []

        for message in messages:
            if isinstance(message, MultiModalMessage):
                role = (
                    message.role.value
                    if hasattr(message.role, "value")
                    else message.role
                )
                content_parts: List[Dict[str, Any]] = []

                if message.content:
                    content_parts.append({"type": "text", "text": message.content})

                if message.is_image():
                    image_url = message.media_url
                    if image_url is None and message.data is not None:
                        base64_data = message.to_base64()
                        if base64_data:
                            image_url = f"data:{message.mime_type};base64,{base64_data}"
                    if image_url:
                        content_parts.append(
                            {"type": "image_url", "image_url": {"url": image_url}}
                        )
                elif message.is_text() and message.data is not None:
                    if isinstance(message.data, bytes):
                        data_text = message.data.decode("utf-8", errors="ignore")
                    else:
                        data_text = message.data
                    if data_text and data_text != message.content:
                        content_parts.append({"type": "text", "text": data_text})

                api_messages.append(
                    {
                        "role": role,
                        "content": content_parts or message.content,
                    }
                )
                continue

            # Build base message dict
            api_message: Dict[str, Any] = {
                "role": message.role.value if hasattr(message.role, "value") else message.role,
                "content": message.content,
            }

            # Handle AssistantMessage with tool calls
            if isinstance(message, AssistantMessage) and message.tool_calls:
                api_message["tool_calls"] = [
                    {
                        "id": tc.call_id,
                        "type": "function",
                        "function": {
                            "name": tc.tool_name,
                            "arguments": json.dumps(tc.parameters)
                            if isinstance(tc.parameters, dict)
                            else tc.parameters,
                        },
                    }
                    for tc in message.tool_calls
                ]

            # Handle ToolMessage (result from tool execution)
            if isinstance(message, ToolMessage):
                api_message["tool_call_id"] = message.tool_call_id

            api_messages.append(api_message)

        return api_messages


# Exception Hierarchy
class RateLimitError(ModelClientError):
    """
    Raised when LLM provider API rate limits are exceeded.

    Indicates temporary unavailability due to rate limiting.
    Recommended action: Exponential backoff retry.
    """

    pass


class AuthenticationError(ModelClientError):
    """
    Raised when LLM provider authentication fails.

    Common causes:
    - Invalid or expired API key
    - Missing authentication credentials
    - IP address not whitelisted
    - API key permissions insufficient for operation
    """

    pass


class InvalidRequestError(ModelClientError):
    """
    Raised when the API request is invalid.

    Common causes:
    - Malformed request format
    - Invalid parameter values
    - Unknown model identifier
    - Tool schema doesn't match provider requirements
    """

    pass


__all__ = [
    "BaseChatCompletionClient",
    "BaseChatCompletionClientConfig",
    "RateLimitError",
    "AuthenticationError",
    "InvalidRequestError",
]

