# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import logging
import sys
from abc import ABC, abstractmethod
from collections.abc import (
    AsyncIterable,
    Awaitable,
    Callable,
    Mapping,
    Sequence,
)
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Generic,
    Literal,
    Protocol,
    TypedDict,
    cast,
    overload,
    runtime_checkable,
)

from ._docstrings import apply_layered_docstring
from ._serialization import SerializationMixin
from ._types import (
    ChatResponse,
    ChatResponseUpdate,
    EmbeddingGenerationOptions,
    EmbeddingInputT,
    EmbeddingT,
    GeneratedEmbeddings,
    Message,
    ResponseStream,
    validate_chat_options,
)

if sys.version_info >= (3, 13):
    from typing import TypeVar  # pragma: no cover
else:
    from typing_extensions import TypeVar  # pragma: no cover


if TYPE_CHECKING:
    from pydantic import BaseModel

    from ._agents import Agent
    from ._compaction import CompactionStrategy, TokenizerProtocol
    from ._middleware import (
        MiddlewareTypes,
    )
    from ._tools import ToolTypes
    from ._types import ChatOptions


InputT = TypeVar("InputT", contravariant=True)

BaseChatClientT = TypeVar("BaseChatClientT", bound="BaseChatClient")

logger = logging.getLogger("agent_framework")


# region SupportsChatGetResponse Protocol

# Contravariant for the Protocol
OptionsContraT = TypeVar(
    "OptionsContraT",
    bound=TypedDict,  # type: ignore[valid-type]
    default="ChatOptions[None]",
    contravariant=True,
)

# Used for the overloads that capture the response model type from options
if TYPE_CHECKING:
    ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel)
else:
    ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=Any)


@runtime_checkable
class SupportsChatGetResponse(Protocol[OptionsContraT]):
    """A protocol for a chat client that can generate responses.

    This protocol defines the interface that all chat clients must implement,
    including methods for generating both streaming and non-streaming responses.

    The generic type parameter TOptions specifies which options TypedDict this
    client accepts, enabling IDE autocomplete and type checking for provider-specific
    options.

    Note:
        Protocols use structural subtyping (duck typing). Classes don't need
        to explicitly inherit from this protocol to be considered compatible.

    Examples:
        .. code-block:: python

            from agent_framework import SupportsChatGetResponse, ChatResponse, Message


            # Any class implementing the required methods is compatible
            class CustomChatClient:
                additional_properties: dict = {}

                def get_response(self, messages, *, stream=False, client_kwargs=None, **kwargs):
                    if stream:
                        from agent_framework import ChatResponseUpdate, ResponseStream

                        async def _stream():
                            yield ChatResponseUpdate()

                        return ResponseStream(_stream())
                    else:

                        async def _response():
                            return ChatResponse(messages=[], response_id="custom")

                        return _response()


            # Verify the instance satisfies the protocol
            client = CustomChatClient()
            assert isinstance(client, SupportsChatGetResponse)
    """

    additional_properties: dict[str, Any]

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[False] = ...,
        options: ChatOptions[ResponseModelBoundT],
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        function_invocation_kwargs: Mapping[str, Any] | None = None,
        client_kwargs: Mapping[str, Any] | None = None,
    ) -> Awaitable[ChatResponse[ResponseModelBoundT]]: ...

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[False] = ...,
        options: OptionsContraT | ChatOptions[None] | None = None,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        function_invocation_kwargs: Mapping[str, Any] | None = None,
        client_kwargs: Mapping[str, Any] | None = None,
    ) -> Awaitable[ChatResponse[Any]]: ...

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[True],
        options: OptionsContraT | ChatOptions[Any] | None = None,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        function_invocation_kwargs: Mapping[str, Any] | None = None,
        client_kwargs: Mapping[str, Any] | None = None,
    ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...

    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: bool = False,
        options: OptionsContraT | ChatOptions[Any] | None = None,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        function_invocation_kwargs: Mapping[str, Any] | None = None,
        client_kwargs: Mapping[str, Any] | None = None,
    ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
        """Send input and return the response.

        Args:
            messages: The sequence of input messages to send.
            stream: Whether to stream the response. Defaults to False.
            options: Chat options as a TypedDict.
            compaction_strategy: Optional per-call compaction override.
            tokenizer: Optional per-call tokenizer override.
            function_invocation_kwargs: Keyword arguments forwarded only to tool invocation layers.
            client_kwargs: Additional client-specific keyword arguments.

        Returns:
            When stream=False: An awaitable ChatResponse from the client.
            When stream=True: A ResponseStream yielding partial updates.

        Raises:
            ValueError: If the input message sequence is ``None``.
        """
        ...


# endregion


# region ChatClientBase

# Covariant for the BaseChatClient
OptionsCoT = TypeVar(
    "OptionsCoT",
    bound=TypedDict,  # type: ignore[valid-type]
    default="ChatOptions[None]",
    covariant=True,
)


class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
    """Abstract base class for chat clients without middleware wrapping.

    This abstract base class provides core functionality for chat client implementations,
    including message preparation and tool normalization, but without middleware,
    telemetry, or function invocation support.

    The generic type parameter TOptions specifies which options TypedDict this client
    accepts. This enables IDE autocomplete and type checking for provider-specific options
    when using the typed overloads of get_response.

    Note:
        BaseChatClient cannot be instantiated directly as it's an abstract base class.
        Subclasses must implement ``_inner_get_response()`` with a stream parameter to handle both
        streaming and non-streaming responses.

        For full-featured clients with middleware, telemetry, and function invocation support,
        use public client classes such as ``OpenAIChatClient`` which compose these layers correctly.

    Examples:
        .. code-block:: python

            from agent_framework import BaseChatClient, ChatResponse, Message
            from collections.abc import AsyncIterable


            class CustomChatClient(BaseChatClient):
                async def _inner_get_response(self, *, messages, stream, options, **kwargs):
                    if stream:
                        # Streaming implementation
                        from agent_framework import ChatResponseUpdate

                        async def _stream():
                            yield ChatResponseUpdate(role="assistant", contents=[{"type": "text", "text": "Hello!"}])

                        return _stream()
                    else:
                        # Non-streaming implementation
                        return ChatResponse(
                            messages=[Message(role="assistant", contents=["Hello!"])],
                            response_id="custom-response",
                        )


            # Create an instance of your custom client
            client = CustomChatClient()

            # Use the client to get responses
            response = await client.get_response([Message(role="user", contents=["Hello, how are you?"])])
            # Or stream responses
            async for update in client.get_response([Message(role="user", contents=["Hello!"])], stream=True):
                print(update)
    """

    OTEL_PROVIDER_NAME: ClassVar[str] = "unknown"
    compaction_strategy: CompactionStrategy | None = None
    tokenizer: TokenizerProtocol | None = None
    DEFAULT_EXCLUDE: ClassVar[set[str]] = {
        "additional_properties",
        "compaction_strategy",
        "tokenizer",
    }
    STORES_BY_DEFAULT: ClassVar[bool] = False
    """Whether this client stores conversation history server-side by default.

    Clients that use server-side storage (e.g., OpenAI Responses API with ``store=True``
    as default, Azure AI Agent sessions) should override this to ``True``.
    When ``True``, the agent skips auto-injecting ``InMemoryHistoryProvider`` unless the
    user explicitly sets ``store=False``.
    """
    # OTEL_PROVIDER_NAME is used for OTel setup, should be overridden in subclasses

    def __init__(
        self,
        *,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        additional_properties: dict[str, Any] | None = None,
    ) -> None:
        """Initialize a BaseChatClient instance.

        Keyword Args:
            compaction_strategy: Optional compaction strategy to apply before model calls.
            tokenizer: Optional tokenizer used by token-aware compaction strategies.
            additional_properties: Additional properties for the client.
        """
        self.additional_properties = additional_properties or {}
        self.compaction_strategy = compaction_strategy
        self.tokenizer = tokenizer
        super().__init__()

    def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
        """Convert the instance to a dictionary.

        Extracts additional_properties fields to the root level.

        Keyword Args:
            exclude: Set of field names to exclude from serialization.
            exclude_none: Whether to exclude None values from the output. Defaults to True.

        Returns:
            Dictionary representation of the instance.
        """
        # Get the base dict from SerializationMixin
        result = super().to_dict(exclude=exclude, exclude_none=exclude_none)

        # Extract additional_properties to root level
        if self.additional_properties:
            result.update(self.additional_properties)

        return result

    async def _validate_options(self, options: Mapping[str, Any]) -> dict[str, Any]:
        """Validate and normalize chat options.

        Subclasses should call this at the start of _inner_get_response to validate options.

        Args:
            options: The raw options dict.

        Returns:
            The validated and normalized options dict.
        """
        return await validate_chat_options(dict(options))

    def _finalize_response_updates(
        self,
        updates: Sequence[ChatResponseUpdate],
        *,
        response_format: Any | None = None,
    ) -> ChatResponse[Any]:
        """Finalize response updates into a single ChatResponse."""
        return ChatResponse.from_updates(
            updates,
            output_format_type=response_format,
        )

    def _build_response_stream(
        self,
        stream: AsyncIterable[ChatResponseUpdate] | Awaitable[AsyncIterable[ChatResponseUpdate]],
        *,
        response_format: Any | None = None,
    ) -> ResponseStream[ChatResponseUpdate, ChatResponse]:
        """Create a ResponseStream with the standard finalizer."""
        return ResponseStream(
            stream,
            finalizer=lambda updates: self._finalize_response_updates(updates, response_format=response_format),
        )

    async def _prepare_messages_for_model_call(
        self,
        messages: Sequence[Message],
        *,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
    ) -> list[Message]:
        prepared_messages = list(messages)
        if compaction_strategy is None:
            if tokenizer is None:
                return prepared_messages
            from ._compaction import annotate_message_groups

            annotate_message_groups(prepared_messages, tokenizer=tokenizer)
            return prepared_messages
        from ._compaction import apply_compaction

        # Compact the caller's list in place when possible. A compaction operation has
        # two halves: exclusion flags (mutated on shared Message objects) and inserted
        # summary messages. Operating on the original list keeps both halves on the list
        # the function-invocation tool loop reuses across iterations; otherwise inserted
        # summaries would be lost on a throwaway copy while exclusions persisted, silently
        # dropping older groups (issue #4991).
        working_messages = messages if isinstance(messages, list) else prepared_messages
        return await apply_compaction(
            working_messages,
            strategy=compaction_strategy,
            tokenizer=tokenizer,
        )

    def _resolve_compaction_overrides(
        self,
        *,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
    ) -> dict[str, Any]:
        current_compaction_strategy = getattr(self, "compaction_strategy", None)
        current_tokenizer = getattr(self, "tokenizer", None)
        ret: dict[str, Any] = {}
        if current_compaction_strategy is not None or compaction_strategy is not None:
            ret["compaction_strategy"] = (
                current_compaction_strategy if compaction_strategy is None else compaction_strategy
            )
        if current_tokenizer is not None or tokenizer is not None:
            ret["tokenizer"] = current_tokenizer if tokenizer is None else tokenizer
        return ret

    # region Internal method to be implemented by derived classes

    @abstractmethod
    def _inner_get_response(
        self,
        *,
        messages: Sequence[Message],
        stream: bool,
        options: Mapping[str, Any],
        **kwargs: Any,
    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
        """Send a chat request to the AI service.

        Subclasses must implement this method to handle both streaming and non-streaming
        responses based on the stream parameter. Implementations should call
        ``await self._validate_options(options)`` at the start to validate options.

        Keyword Args:
            messages: The prepared chat messages to send.
            stream: Whether to stream the response.
            options: The options dict for the request (call _validate_options first).
            kwargs: Any additional keyword arguments.

        Returns:
            When stream=False: An Awaitable ChatResponse from the model.
            When stream=True: A ResponseStream of ChatResponseUpdate instances.
        """

    # region Public method

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[False] = ...,
        options: ChatOptions[ResponseModelBoundT],
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        function_invocation_kwargs: Mapping[str, Any] | None = None,
        client_kwargs: Mapping[str, Any] | None = None,
    ) -> Awaitable[ChatResponse[ResponseModelBoundT]]: ...

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[False] = ...,
        options: OptionsCoT | ChatOptions[None] | None = None,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        function_invocation_kwargs: Mapping[str, Any] | None = None,
        client_kwargs: Mapping[str, Any] | None = None,
    ) -> Awaitable[ChatResponse[Any]]: ...

    @overload
    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: Literal[True],
        options: OptionsCoT | ChatOptions[Any] | None = None,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        function_invocation_kwargs: Mapping[str, Any] | None = None,
        client_kwargs: Mapping[str, Any] | None = None,
    ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...

    def get_response(
        self,
        messages: Sequence[Message],
        *,
        stream: bool = False,
        options: OptionsCoT | ChatOptions[Any] | None = None,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        function_invocation_kwargs: Mapping[str, Any] | None = None,
        client_kwargs: Mapping[str, Any] | None = None,
    ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
        """Get a response from a chat client.

        Args:
            messages: The message or messages to send to the model.
            stream: Whether to stream the response. Defaults to False.
            options: Chat options as a TypedDict.
            compaction_strategy: Optional per-call override for in-run compaction.
                When omitted, the client-level default is used.
            tokenizer: Optional per-call tokenizer override. When omitted, the
                client-level default is used.
            function_invocation_kwargs: Keyword arguments forwarded only to tool invocation layers.
            client_kwargs: Additional client-specific keyword arguments forwarded to
                ``_inner_get_response()``.

        Returns:
            When streaming a response stream of ChatResponseUpdates, otherwise an Awaitable ChatResponse.
        """
        compaction_overrides = self._resolve_compaction_overrides(
            compaction_strategy=compaction_strategy,
            tokenizer=tokenizer,
        )
        merged_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {}

        if not compaction_overrides:
            return self._inner_get_response(
                messages=messages,
                stream=stream,
                options=options or {},
                **merged_client_kwargs,
            )

        if stream:

            async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
                prepared_messages = await self._prepare_messages_for_model_call(
                    messages,
                    **compaction_overrides,
                )
                stream_response = self._inner_get_response(
                    messages=prepared_messages,
                    stream=True,
                    options=options or {},
                    **merged_client_kwargs,
                )
                if isinstance(stream_response, ResponseStream):
                    return stream_response  # type: ignore[reportUnknownVariableType]
                awaited_stream_response = await stream_response
                if isinstance(awaited_stream_response, ResponseStream):
                    return awaited_stream_response
                raise ValueError("Streaming responses must return a ResponseStream.")

            return ResponseStream.from_awaitable(_get_stream())  # type: ignore[reportUnknownVariableType]

        async def _get_response() -> ChatResponse[Any]:
            prepared_messages = await self._prepare_messages_for_model_call(
                messages,
                **compaction_overrides,
            )
            return await self._inner_get_response(
                messages=prepared_messages,
                stream=False,
                options=options or {},
                **merged_client_kwargs,
            )

        return _get_response()

    def service_url(self) -> str:
        """Get the URL of the service.

        Override this in the subclass to return the proper URL.
        If the service does not have a URL, return None.

        Returns:
            The service URL or 'Unknown' if not implemented.
        """
        return "Unknown"

    def as_agent(
        self,
        *,
        id: str | None = None,
        name: str | None = None,
        description: str | None = None,
        instructions: str | None = None,
        tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
        default_options: OptionsCoT | Mapping[str, Any] | None = None,
        context_providers: Sequence[Any] | None = None,
        middleware: Sequence[MiddlewareTypes] | None = None,
        require_per_service_call_history_persistence: bool = False,
        compaction_strategy: CompactionStrategy | None = None,
        tokenizer: TokenizerProtocol | None = None,
        additional_properties: Mapping[str, Any] | None = None,
    ) -> Agent[OptionsCoT]:
        """Create an Agent with this client.

        This is a convenience method that creates an Agent instance with this
        chat client already configured.

        Keyword Args:
            id: The unique identifier for the agent. Will be created automatically if not provided.
            name: The name of the agent.
            description: A brief description of the agent's purpose.
            instructions: Optional instructions for the agent.
                These will be put into the messages sent to the chat client service as a system message.
            tools: The tools to use for the request.
            default_options: A TypedDict containing chat options. When using a typed client like
                ``OpenAIChatClient``, this enables IDE autocomplete for provider-specific options
                including temperature, max_tokens, model, tool_choice, and more.
                Note: response_format typing does not flow into run outputs when set via default_options,
                and dict literals are accepted without specialized option typing.
            context_providers: Context providers to include during agent invocation.
            middleware: List of middleware to intercept agent and function invocations.
            require_per_service_call_history_persistence: When enabled (and a HistoryProvider is
                present), the provider always persists history after each model call. If the
                client does not store history server-side, history providers are also loaded and
                injected around each model call; if it does, provider loading is skipped and the
                service-managed conversation is the source of truth (persistence still happens
                after each model call). When no HistoryProvider is present, this flag has no
                effect (no middleware is installed and nothing is persisted).
            compaction_strategy: Optional agent-level compaction override. When omitted,
                client-level compaction defaults remain in effect for each call.
            tokenizer: Optional agent-level tokenizer override. When omitted,
                client-level tokenizer defaults remain in effect for each call.
            additional_properties: Additional properties stored on the created agent.

        Returns:
            An Agent instance configured with this chat client.

        Examples:
            .. code-block:: python

                from agent_framework.openai import OpenAIChatClient

                # Create a client
                client = OpenAIChatClient(model="gpt-4")

                # Create an agent using the convenience method
                agent = client.as_agent(
                    name="assistant",
                    instructions="You are a helpful assistant.",
                    default_options={"temperature": 0.7, "max_tokens": 500},
                )

                # Run the agent
                response = await agent.run("Hello!")
        """
        from ._agents import Agent

        agent_kwargs: dict[str, Any] = {
            "client": self,
            "id": id,
            "name": name,
            "description": description,
            "instructions": instructions,
            "tools": tools,
            "default_options": cast(Any, default_options),
            "context_providers": context_providers,
            "middleware": middleware,
            "require_per_service_call_history_persistence": require_per_service_call_history_persistence,
            "compaction_strategy": compaction_strategy,
            "tokenizer": tokenizer,
            "additional_properties": dict(additional_properties) if additional_properties is not None else None,
        }

        return Agent(**agent_kwargs)


# endregion


# region Tool Support Protocols


@runtime_checkable
class SupportsCodeInterpreterTool(Protocol):
    """Protocol for clients that support code interpreter tools.

    This protocol enables runtime checking to determine if a client
    supports code interpreter functionality.

    Examples:
        .. code-block:: python

            from agent_framework import Agent, SupportsCodeInterpreterTool

            if isinstance(client, SupportsCodeInterpreterTool):
                tool = client.get_code_interpreter_tool()
                agent = Agent(client, tools=[tool])
    """

    @staticmethod
    def get_code_interpreter_tool(**kwargs: Any) -> Any:
        """Create a code interpreter tool configuration.

        Keyword Args:
            **kwargs: Provider-specific configuration options.

        Returns:
            A tool configuration ready to pass to Agent.
        """
        ...


@runtime_checkable
class SupportsWebSearchTool(Protocol):
    """Protocol for clients that support web search tools.

    This protocol enables runtime checking to determine if a client
    supports web search functionality.

    Examples:
        .. code-block:: python

            from agent_framework import Agent, SupportsWebSearchTool

            if isinstance(client, SupportsWebSearchTool):
                tool = client.get_web_search_tool()
                agent = Agent(client, tools=[tool])
    """

    @staticmethod
    def get_web_search_tool(**kwargs: Any) -> Any:
        """Create a web search tool configuration.

        Keyword Args:
            **kwargs: Provider-specific configuration options.

        Returns:
            A tool configuration ready to pass to Agent.
        """
        ...


@runtime_checkable
class SupportsImageGenerationTool(Protocol):
    """Protocol for clients that support image generation tools.

    This protocol enables runtime checking to determine if a client
    supports image generation functionality.

    Examples:
        .. code-block:: python

            from agent_framework import Agent, SupportsImageGenerationTool

            if isinstance(client, SupportsImageGenerationTool):
                tool = client.get_image_generation_tool()
                agent = Agent(client, tools=[tool])
    """

    @staticmethod
    def get_image_generation_tool(**kwargs: Any) -> Any:
        """Create an image generation tool configuration.

        Keyword Args:
            **kwargs: Provider-specific configuration options.

        Returns:
            A tool configuration ready to pass to Agent.
        """
        ...


@runtime_checkable
class SupportsMCPTool(Protocol):
    """Protocol for clients that support MCP (Model Context Protocol) tools.

    This protocol enables runtime checking to determine if a client
    supports MCP server connections.

    Examples:
        .. code-block:: python

            from agent_framework import Agent, SupportsMCPTool

            if isinstance(client, SupportsMCPTool):
                tool = client.get_mcp_tool(name="my_mcp", url="https://...")
                agent = Agent(client, tools=[tool])
    """

    @staticmethod
    def get_mcp_tool(**kwargs: Any) -> Any:
        """Create an MCP tool configuration.

        Keyword Args:
            **kwargs: Provider-specific configuration options including
                name and url for the MCP server.

        Returns:
            A tool configuration ready to pass to Agent.
        """
        ...


@runtime_checkable
class SupportsFileSearchTool(Protocol):
    """Protocol for clients that support file search tools.

    This protocol enables runtime checking to determine if a client
    supports file search functionality with vector stores.

    Examples:
        .. code-block:: python

            from agent_framework import Agent, SupportsFileSearchTool

            if isinstance(client, SupportsFileSearchTool):
                tool = client.get_file_search_tool(vector_store_ids=["vs_123"])
                agent = Agent(client, tools=[tool])
    """

    @staticmethod
    def get_file_search_tool(**kwargs: Any) -> Any:
        """Create a file search tool configuration.

        Keyword Args:
            **kwargs: Provider-specific configuration options.

        Returns:
            A tool configuration ready to pass to Agent.
        """
        ...


@runtime_checkable
class SupportsShellTool(Protocol):
    """Protocol for clients that support shell tools.

    This protocol enables runtime checking to determine if a client
    supports executing shell commands.

    Examples:
        .. code-block:: python

            from agent_framework import Agent, SupportsShellTool
            from agent_framework_tools.shell import LocalShellTool


            async def build_agent(client):
                if isinstance(client, SupportsShellTool):
                    async with LocalShellTool() as shell:
                        tool = client.get_shell_tool(func=shell.as_function())
                        return Agent(client, tools=[tool])
    """

    @staticmethod
    def get_shell_tool(**kwargs: Any) -> Any:
        """Create a shell tool configuration.

        Keyword Args:
            **kwargs: Provider-specific configuration options.

        Returns:
            A tool configuration ready to pass to Agent.
        """
        ...


# endregion


# region SupportsGetEmbeddings Protocol

# TypeVars for the Protocol
EmbeddingInputContraT = TypeVar(
    "EmbeddingInputContraT",
    default="str",
    contravariant=True,
)
EmbeddingProtocolOptionsT = TypeVar(
    "EmbeddingProtocolOptionsT",
    bound=TypedDict,  # type: ignore[valid-type]
    default="EmbeddingGenerationOptions",
)


@runtime_checkable
class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingT, EmbeddingProtocolOptionsT]):
    """Protocol for an embedding client that can generate embeddings.

    This protocol enables duck-typing for embedding generation. Any class that
    implements ``get_embeddings`` with a compatible signature satisfies this protocol.

    Generic over the input type (defaults to ``str``), output embedding type
    (defaults to ``list[float]``), and options type.

    Examples:
        .. code-block:: python

            from agent_framework import SupportsGetEmbeddings


            async def use_embeddings(client: SupportsGetEmbeddings) -> None:
                result = await client.get_embeddings(["Hello, world!"])
                for embedding in result:
                    print(embedding.vector)
    """

    additional_properties: dict[str, Any]

    def get_embeddings(
        self,
        values: Sequence[EmbeddingInputContraT],
        *,
        options: EmbeddingProtocolOptionsT | None = None,
    ) -> Awaitable[GeneratedEmbeddings[EmbeddingT, EmbeddingProtocolOptionsT]]:
        """Generate embeddings for the given values.

        Args:
            values: The values to generate embeddings for.
            options: Optional embedding generation options.

        Returns:
            Generated embeddings with metadata.
        """
        ...


# endregion


# region BaseEmbeddingClient

# Covariant for the BaseEmbeddingClient
EmbeddingOptionsT = TypeVar(
    "EmbeddingOptionsT",
    bound=TypedDict,  # type: ignore[valid-type]
    default="EmbeddingGenerationOptions",
    covariant=True,
)


class BaseEmbeddingClient(SerializationMixin, ABC, Generic[EmbeddingInputT, EmbeddingT, EmbeddingOptionsT]):
    """Abstract base class for embedding clients.

    Subclasses implement ``get_embeddings`` to provide the actual
    embedding generation logic.

    Generic over the input type (defaults to ``str``), output embedding type
    (defaults to ``list[float]``), and options type.

    Examples:
        .. code-block:: python

            from agent_framework import BaseEmbeddingClient, Embedding, GeneratedEmbeddings
            from collections.abc import Sequence


            class CustomEmbeddingClient(BaseEmbeddingClient):
                async def get_embeddings(self, values, *, options=None):
                    return GeneratedEmbeddings([Embedding(vector=[0.1, 0.2, 0.3]) for _ in values])
    """

    OTEL_PROVIDER_NAME: ClassVar[str] = "unknown"
    DEFAULT_EXCLUDE: ClassVar[set[str]] = {"additional_properties"}

    def __init__(
        self,
        *,
        additional_properties: dict[str, Any] | None = None,
    ) -> None:
        """Initialize a BaseEmbeddingClient instance.

        Args:
            additional_properties: Additional properties to pass to the client.
        """
        self.additional_properties = additional_properties or {}
        super().__init__()

    @abstractmethod
    async def get_embeddings(
        self,
        values: Sequence[EmbeddingInputT],
        *,
        options: EmbeddingOptionsT | None = None,
    ) -> GeneratedEmbeddings[EmbeddingT, EmbeddingOptionsT]:
        """Generate embeddings for the given values.

        Args:
            values: The values to generate embeddings for.
            options: Optional embedding generation options.

        Returns:
            Generated embeddings with metadata.
        """
        ...


# endregion


def _apply_get_response_docstrings() -> None:
    """Align layered chat-client docstrings with the lowest public implementation."""
    try:
        from ._middleware import ChatMiddlewareLayer
    except ImportError as exc:
        if exc.name == "agent_framework._middleware" and "partially initialized module" in str(exc):
            return
        raise

    from ._tools import FunctionInvocationLayer
    from .observability import ChatTelemetryLayer

    apply_layered_docstring(ChatTelemetryLayer.get_response, BaseChatClient.get_response)
    apply_layered_docstring(FunctionInvocationLayer.get_response, ChatTelemetryLayer.get_response)
    apply_layered_docstring(
        ChatMiddlewareLayer.get_response,
        FunctionInvocationLayer.get_response,
        extra_keyword_args={
            "middleware": """
                Optional per-call chat and function middleware.
                This compatibility keyword argument is merged with any ``client_kwargs["middleware"]`` value
                before the request is executed.
            """,
        },
    )


_apply_get_response_docstrings()
