from __future__ import annotations as _annotations

from abc import ABC
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import datetime
from typing import Annotated, Any, Literal, Union

import pydantic
from pydantic_core import core_schema
from typing_extensions import TypedDict, deprecated

__all__ = (
    'AbstractNativeTool',
    'WebSearchTool',
    'WebSearchUserLocation',
    'XSearchTool',
    'CodeExecutionTool',
    'WebFetchTool',
    'UrlContextTool',
    'ImageGenerationModelName',
    'ImageGenerationTool',
    'ImageAspectRatio',
    'MemoryTool',
    'MCPServerTool',
    'FileSearchTool',
    'NATIVE_TOOL_TYPES',
    'DEPRECATED_NATIVE_TOOLS',
    'SUPPORTED_NATIVE_TOOLS',
    'NATIVE_TOOLS_REQUIRING_CONFIG',
)

NATIVE_TOOL_TYPES: dict[str, type[AbstractNativeTool]] = {}
"""Registry of all native tool types, keyed by their kind string.

This dict is populated automatically via `__init_subclass__` when tool classes are defined.
"""

ImageAspectRatio = Literal['21:9', '16:9', '4:3', '3:2', '1:1', '9:16', '3:4', '2:3', '5:4', '4:5']
"""Supported aspect ratios for image generation tools."""

ImageGenerationModelName = Literal['gpt-image-2', 'gpt-image-1.5', 'gpt-image-1', 'gpt-image-1-mini'] | str
"""Known OpenAI image generation model names, or another OpenAI image model ID."""


@dataclass(kw_only=True)
class AbstractNativeTool(ABC):
    """A native tool that can be used by an agent.

    This class is abstract and cannot be instantiated directly.

    The native tools are passed to the model as part of the `ModelRequestParameters`.
    """

    kind: str = 'unknown_native_tool'
    """Native tool identifier, this should be available on all native tools as a discriminator."""

    optional: bool = False
    """Whether this instance is a best-effort upgrade rather than a hard requirement.

    When `True`, the instance is silently dropped from the request on a model that doesn't
    support it natively, instead of raising when no local fallback is provided. Use for
    native tools where a fallback path exists (e.g. a local function tool that takes over when
    the native one isn't available). When `False` (the default), the request errors on
    models that can't honor the native tool — the user explicitly asked for it, so fail loudly
    rather than silently substituting different behavior.
    """

    @property
    def unique_id(self) -> str:
        """A unique identifier for the native tool.

        If multiple instances of the same native tool can be passed to the model, subclasses should override this property to allow them to be distinguished.
        """
        return self.kind

    @property
    def label(self) -> str:
        """Human-readable label for UI display.

        Subclasses should override this to provide a meaningful label.
        """
        return self.kind.replace('_', ' ').title()

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        NATIVE_TOOL_TYPES[cls.kind] = cls

    @classmethod
    def __get_pydantic_core_schema__(
        cls, _source_type: Any, handler: pydantic.GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        if cls is not AbstractNativeTool:
            return handler(cls)

        tools = NATIVE_TOOL_TYPES.values()
        if len(tools) == 1:  # pragma: no cover
            tools_type = next(iter(tools))
        else:
            tools_annotated = [Annotated[tool, pydantic.Tag(tool.kind)] for tool in tools]
            tools_type = Annotated[Union[tuple(tools_annotated)], pydantic.Discriminator(_tool_discriminator)]  # noqa: UP007

        return handler(tools_type)


@dataclass(kw_only=True)
class WebSearchTool(AbstractNativeTool):
    """A native tool that allows your agent to search the web for information.

    The parameters that PydanticAI passes depend on the model, as some parameters may not be supported by certain models.

    Supported by:

    * Anthropic
    * OpenAI Responses
    * Groq
    * Google
    * xAI
    * OpenRouter
    """

    search_context_size: Literal['low', 'medium', 'high'] = 'medium'
    """The `search_context_size` parameter controls how much context is retrieved from the web to help the tool formulate a response.

    Supported by:

    * OpenAI Responses
    * OpenRouter
    """

    user_location: WebSearchUserLocation | None = None
    """The `user_location` parameter allows you to localize search results based on a user's location.

    Supported by:

    * Anthropic
    * OpenAI Responses
    """

    blocked_domains: list[str] | None = None
    """If provided, these domains will never appear in results.

    With Anthropic, you can only use one of `blocked_domains` or `allowed_domains`, not both.

    Supported by:

    * Anthropic, see <https://docs.anthropic.com/en/docs/build-with-claude/tool-use/web-search-tool#domain-filtering>
    * Groq, see <https://console.groq.com/docs/agentic-tooling#search-settings>
    * xAI, see <https://docs.x.ai/docs/guides/tools/search-tools#web-search-parameters>
    """

    allowed_domains: list[str] | None = None
    """If provided, only these domains will be included in results.

    With Anthropic, you can only use one of `blocked_domains` or `allowed_domains`, not both.

    Supported by:

    * Anthropic, see <https://docs.anthropic.com/en/docs/build-with-claude/tool-use/web-search-tool#domain-filtering>
    * Groq, see <https://console.groq.com/docs/agentic-tooling#search-settings>
    * OpenAI Responses, see <https://platform.openai.com/docs/guides/tools-web-search>
    * xAI, see <https://docs.x.ai/docs/guides/tools/search-tools#web-search-parameters>
    """

    max_uses: int | None = None
    """If provided, the tool will stop searching the web after the given number of uses.

    Supported by:

    * Anthropic
    """

    kind: str = 'web_search'
    """The kind of tool."""


class WebSearchUserLocation(TypedDict, total=False):
    """Allows you to localize search results based on a user's location.

    Supported by:

    * Anthropic
    * OpenAI Responses
    """

    city: str
    """The city where the user is located."""

    country: str
    """The country where the user is located. For OpenAI, this must be a 2-letter country code (e.g., 'US', 'GB')."""

    region: str
    """The region or state where the user is located."""

    timezone: str
    """The timezone of the user's location."""


@dataclass(kw_only=True)
class XSearchTool(AbstractNativeTool):
    """A native tool that allows your agent to search X/Twitter for posts and content.

    See <https://docs.x.ai/developers/tools/x-search> for more details.

    When used via the [`XSearch`][pydantic_ai.capabilities.XSearch] capability with a
    `fallback_model` set, this tool also works with non-xAI models by delegating to a
    subagent running the specified xAI model.

    Supported by:

    * xAI
    """

    allowed_x_handles: list[str] | None = None
    """If provided, only posts from these X handles will be included (max 10).

    Supported by:

    * xAI, see <https://docs.x.ai/developers/tools/x-search>
    """

    excluded_x_handles: list[str] | None = None
    """If provided, posts from these X handles will be excluded (max 10).

    Supported by:

    * xAI, see <https://docs.x.ai/developers/tools/x-search>
    """

    from_date: datetime | None = None
    """If provided, only posts created on or after this datetime will be included.

    Naive datetimes are interpreted as UTC by the xAI API.

    Supported by:

    * xAI, see <https://docs.x.ai/developers/tools/x-search>
    """

    to_date: datetime | None = None
    """If provided, only posts created on or before this datetime will be included.

    Naive datetimes are interpreted as UTC by the xAI API.

    Supported by:

    * xAI, see <https://docs.x.ai/developers/tools/x-search>
    """

    enable_image_understanding: bool = False
    """Enable image analysis from X posts.

    Supported by:

    * xAI, see <https://docs.x.ai/developers/tools/x-search>
    """

    enable_video_understanding: bool = False
    """Enable video analysis from X content.

    Supported by:

    * xAI, see <https://docs.x.ai/developers/tools/x-search>
    """

    include_output: bool = False
    """Include raw X search results in the response as
    [`NativeToolReturnPart`][pydantic_ai.messages.NativeToolReturnPart].

    Without this, the model uses the search results internally but only returns
    its text summary. Enabling it gives programmatic access to searched posts,
    sources, and metadata.

    Can also be set via
    [`XaiModelSettings.xai_include_x_search_output`][pydantic_ai.models.xai.XaiModelSettings.xai_include_x_search_output].

    Supported by:

    * xAI, see <https://docs.x.ai/developers/tools/x-search>
    """

    kind: str = 'x_search'
    """The kind of tool."""

    def __post_init__(self) -> None:
        if self.allowed_x_handles is not None and self.excluded_x_handles is not None:
            raise ValueError('Cannot specify both allowed_x_handles and excluded_x_handles')
        if self.allowed_x_handles and len(self.allowed_x_handles) > 10:
            raise ValueError('allowed_x_handles cannot contain more than 10 handles')
        if self.excluded_x_handles and len(self.excluded_x_handles) > 10:
            raise ValueError('excluded_x_handles cannot contain more than 10 handles')


@dataclass(kw_only=True)
class CodeExecutionTool(AbstractNativeTool):
    """A native tool that allows your agent to execute code.

    Supported by:

    * Anthropic
    * OpenAI Responses
    * Google
    * Bedrock (Nova2.0)
    * xAI
    """

    kind: str = 'code_execution'
    """The kind of tool."""


@dataclass(kw_only=True)
class WebFetchTool(AbstractNativeTool):
    """Allows your agent to access contents from URLs.

    The parameters that PydanticAI passes depend on the model, as some parameters may not be supported by certain models.

    Supported by:

    * Anthropic
    * Google
    """

    max_uses: int | None = None
    """If provided, the tool will stop fetching URLs after the given number of uses.

    Supported by:

    * Anthropic
    """

    allowed_domains: list[str] | None = None
    """If provided, only these domains will be fetched.

    With Anthropic, you can only use one of `blocked_domains` or `allowed_domains`, not both.

    Supported by:

    * Anthropic, see <https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/web-fetch-tool#domain-filtering>
    """

    blocked_domains: list[str] | None = None
    """If provided, these domains will never be fetched.

    With Anthropic, you can only use one of `blocked_domains` or `allowed_domains`, not both.

    Supported by:

    * Anthropic, see <https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/web-fetch-tool#domain-filtering>
    """

    enable_citations: bool = False
    """If True, enables citations for fetched content.

    Supported by:

    * Anthropic
    """

    max_content_tokens: int | None = None
    """Maximum content length in tokens for fetched content.

    Supported by:

    * Anthropic
    """

    kind: str = 'web_fetch'
    """The kind of tool."""


@deprecated('Use `WebFetchTool` instead.')
@dataclass(kw_only=True)
class UrlContextTool(WebFetchTool):
    """Deprecated alias for WebFetchTool. Use WebFetchTool instead.

    Overrides kind to 'url_context' so old serialized payloads with {"kind": "url_context", ...}
    can be deserialized to UrlContextTool for backward compatibility.
    """

    kind: str = 'url_context'
    """The kind of tool (deprecated value for backward compatibility)."""


@dataclass(kw_only=True)
class ImageGenerationTool(AbstractNativeTool):
    """A native tool that allows your agent to generate images.

    Supported by:

    * OpenAI Responses
    * Google
    """

    action: Literal['generate', 'edit', 'auto'] = 'auto'
    """Whether to generate a new image or edit an existing image.

    Supported by:

    * OpenAI Responses. Default: 'auto'.
    """

    background: Literal['transparent', 'opaque', 'auto'] = 'auto'
    """Background type for the generated image.

    Supported by:

    * OpenAI Responses. 'transparent' is only supported for 'png' and 'webp' output formats.
    """

    input_fidelity: Literal['high', 'low'] | None = None
    """
    Control how much effort the model will exert to match the style and features,
    especially facial features, of input images.

    Supported by:

    * OpenAI Responses. Default: 'low'.
    """

    moderation: Literal['auto', 'low'] = 'auto'
    """Moderation level for the generated image.

    Supported by:

    * OpenAI Responses
    """

    model: ImageGenerationModelName | None = None
    """The image generation model to use.

    Supported by:

    * OpenAI Responses. Defaults to the provider's image generation model selection.
      Known image generation models include `gpt-image-2`, `gpt-image-1.5`,
      `gpt-image-1`, and `gpt-image-1-mini`.

    This selects the underlying image generation model used by the tool; it does
    not change the agent's conversational model.
    """

    output_compression: int | None = None
    """Compression level for the output image.

    Supported by:

    * OpenAI Responses. Only supported for 'jpeg' and 'webp' output formats. Default: 100.
    * Google (Vertex AI only). Only supported for 'jpeg' output format. Default: 75.
      Setting this will default `output_format` to 'jpeg' if not specified.
    """

    output_format: Literal['png', 'webp', 'jpeg'] | None = None
    """The output format of the generated image.

    Supported by:

    * OpenAI Responses. Default: 'png'.
    * Google (Vertex AI only). Default: 'png', or 'jpeg' if `output_compression` is set.
    """

    partial_images: int = 0
    """
    Number of partial images to generate in streaming mode.

    Supported by:

    * OpenAI Responses. Supports 0 to 3.
    """

    quality: Literal['low', 'medium', 'high', 'auto'] = 'auto'
    """The quality of the generated image.

    Supported by:

    * OpenAI Responses
    """

    size: Literal['auto', '1024x1024', '1024x1536', '1536x1024', '512', '1K', '2K', '4K'] | None = None
    """The size of the generated image.

    * OpenAI Responses: 'auto' (default: model selects the size based on the prompt), '1024x1024', '1024x1536', '1536x1024'
    * Google (Gemini 3 Pro Image and later): '512' (Gemini 3.1 Flash Image only), '1K' (default), '2K', '4K'
    """

    aspect_ratio: ImageAspectRatio | None = None
    """The aspect ratio to use for generated images.

    Supported by:

    * Google image-generation models (Gemini)
    * OpenAI Responses (maps '1:1', '2:3', and '3:2' to supported sizes)
    """

    kind: str = 'image_generation'
    """The kind of tool."""


@dataclass(kw_only=True)
class MemoryTool(AbstractNativeTool):
    """A native tool that allows your agent to use memory.

    Supported by:

    * Anthropic
    """

    kind: str = 'memory'
    """The kind of tool."""


@dataclass(kw_only=True)
class MCPServerTool(AbstractNativeTool):
    """A native tool that allows your agent to use MCP servers.

    Supported by:

    * OpenAI Responses
    * Anthropic
    * xAI
    """

    id: str
    """A unique identifier for the MCP server."""

    url: str
    """The URL of the MCP server to use.

    For OpenAI Responses, it is possible to use `connector_id` by providing it as `x-openai-connector:<connector_id>`.
    """

    authorization_token: str | None = None
    """Authorization header to use when making requests to the MCP server.

    Supported by:

    * OpenAI Responses
    * Anthropic
    * xAI
    """

    description: str | None = None
    """A description of the MCP server.

    Supported by:

    * OpenAI Responses
    * xAI
    """

    allowed_tools: list[str] | None = None
    """A list of tools that the MCP server can use.

    Supported by:

    * OpenAI Responses
    * Anthropic
    * xAI
    """

    headers: dict[str, str] | None = None
    """Optional HTTP headers to send to the MCP server.

    Use for authentication or other purposes.

    Supported by:

    * OpenAI Responses
    * xAI
    """

    kind: str = 'mcp_server'

    @property
    def unique_id(self) -> str:
        return ':'.join([self.kind, self.id])

    @property
    def label(self) -> str:
        return f'MCP: {self.id}'


@dataclass(kw_only=True)
class FileSearchTool(AbstractNativeTool):
    """A native tool that allows your agent to search through uploaded files using vector search.

    This tool provides a fully managed Retrieval-Augmented Generation (RAG) system that handles
    file storage, chunking, embedding generation, and context injection into prompts.

    Supported by:

    * OpenAI Responses
    * Google (Gemini)
    * xAI (mapped to collections search)
    """

    file_store_ids: Sequence[str]
    """The file store IDs to search through.

    For OpenAI, these are the IDs of vector stores created via the OpenAI API.
    For Google, these are file search store names that have been uploaded and processed via the Gemini Files API.
    For xAI, these are collection IDs for the xAI collections search tool.
    """

    kind: str = 'file_search'
    """The kind of tool."""


# Imported after the base class is defined — `_tool_search.py` subclasses
# `AbstractNativeTool`, so the import has to follow. Loading the submodule registers
# `ToolSearchTool` in `NATIVE_TOOL_TYPES` via `__init_subclass__`. `ToolSearchTool` is
# framework-internal (constructed exclusively by the
# [`ToolSearch`][pydantic_ai.capabilities.ToolSearch] capability) and intentionally not
# re-exported here; user-facing strategy types live on `pydantic_ai.capabilities`.
from . import _tool_search as _tool_search  # noqa: E402


def _tool_discriminator(tool_data: dict[str, Any] | AbstractNativeTool) -> str:
    if isinstance(tool_data, dict):
        return tool_data.get('kind', AbstractNativeTool.kind)
    else:
        return tool_data.kind


DEPRECATED_NATIVE_TOOLS: frozenset[type[AbstractNativeTool]] = frozenset({UrlContextTool})  # pyright: ignore[reportDeprecated]
"""Set of deprecated native tool IDs that should not be offered in new UIs."""

SUPPORTED_NATIVE_TOOLS = frozenset(cls for cls in NATIVE_TOOL_TYPES.values() if cls not in DEPRECATED_NATIVE_TOOLS)
"""Get the set of all native tool types (excluding deprecated tools)."""

NATIVE_TOOLS_REQUIRING_CONFIG: frozenset[type[AbstractNativeTool]] = frozenset(
    {FileSearchTool, MCPServerTool, MemoryTool, _tool_search.ToolSearchTool}
)
