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

from pydantic_ai.messages import UploadedFile

__all__ = (
    'AbstractNativeTool',
    'WebSearchTool',
    'WebSearchUserLocation',
    'XSearchTool',
    'CodeExecutionTool',
    'WebFetchTool',
    'ImageGenerationModelName',
    'ImageGenerationTool',
    'ImageAspectRatio',
    'ImageSize',
    'MemoryTool',
    'MCPServerTool',
    'FileSearchTool',
    'AdvisorModelName',
    'AdvisorTool',
    'NATIVE_TOOL_TYPES',
    '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.

The direct image API takes the wider
[`ImageGenerationAspectRatio`][pydantic_ai.images.ImageGenerationAspectRatio], of which these ten are a subset.
"""

ImageSize = Literal['auto', '1024x1024', '1024x1536', '1536x1024', '512', '1K', '2K', '4K']
"""Supported sizes 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."""

AdvisorModelName = (
    Literal[
        'claude-fable-5-1',
        'claude-fable-5',
        'claude-mythos-5-1',
        'claude-mythos-5',
        'claude-opus-5',
        'claude-opus-4-8',
        'claude-opus-4-7',
        'claude-opus-4-6',
        'claude-sonnet-4-6',
    ]
    | str
)
"""Known Anthropic advisor model names, or any other model ID string.

These are the models Anthropic currently accepts as the *advisor* — the stronger model an
executor consults mid-generation. The executor/advisor pairing is validated by the API, not here.
The literals are Anthropic model IDs; on OpenRouter, pass a catalog slug string instead
(e.g. `anthropic/claude-opus-4.8` or the `~anthropic/claude-opus-latest` alias).
"""


@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.
    OpenRouter uses its Beta web-search server tool, which lets the model decide whether to search and how often.
    It accepts the portable settings below, though their effect depends on OpenRouter's selected
    search engine and downstream provider.

    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
    * xAI, see <https://docs.x.ai/docs/guides/tools/search-tools#web-search-parameters>
    * OpenRouter, see <https://openrouter.ai/docs/guides/features/server-tools/web-search#configuration>
    """

    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>
    * OpenAI Responses, see <https://developers.openai.com/api/docs/guides/tools-web-search#domain-filtering>
    * xAI, see <https://docs.x.ai/docs/guides/tools/search-tools#web-search-parameters>
    * OpenRouter, see <https://openrouter.ai/docs/guides/features/server-tools/web-search#configuration>
    """

    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>
    * OpenRouter, see <https://openrouter.ai/docs/guides/features/server-tools/web-search#configuration>
    """

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

    For OpenRouter, this limit is enforced with a non-native search engine or Anthropic's native
    search. Other native providers ignore it.

    Supported by:

    * Anthropic
    * OpenRouter
    """

    external_web_access: bool | None = None
    """Whether the hosted web search tool may fetch live web content.

    If `False`, the tool uses only cached or indexed results. If `None`, the parameter is omitted and the provider
    default is used. OpenAI currently defaults to `True`.

    Supported by:

    * OpenAI Responses `web_search` tool, see <https://developers.openai.com/api/docs/guides/tools-web-search#live-internet-access>

    OpenAI's legacy `web_search_preview` tool ignores this parameter.
    """

    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
    * xAI
    * OpenRouter
    """

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

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

    For OpenAI and xAI, this must be a 2-letter ISO 3166-1 alpha-2 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_subagent_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 20).

    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 20).

    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) > 20:
            raise ValueError('allowed_x_handles cannot contain more than 20 handles')
        if self.excluded_x_handles and len(self.excluded_x_handles) > 20:
            raise ValueError('excluded_x_handles cannot contain more than 20 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
    """

    files: list[UploadedFile] | None = None
    """Uploaded files to make available in the code execution environment.

    Only files matching the model provider are used; files from other providers are ignored.

    Supported by:

    * Anthropic
    * OpenAI Responses
    """

    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."""


@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: ImageSize | 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.
    """

    max_num_results: int | None = None
    """The maximum number of results to return.

    Supported by:

    * xAI (mapped to collections search `limit`, defaults to 10 server-side)
    """

    instructions: str | None = None
    """Optional instructions that guide how the collections search results are interpreted and ranked.

    Supported by:

    * xAI
    """

    retrieval_mode: Literal['hybrid', 'semantic', 'keyword'] | None = None
    """The retrieval strategy for the search.

    Supported by:

    * xAI (defaults to `hybrid` server-side)
    """

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


@dataclass(kw_only=True)
class AdvisorTool(AbstractNativeTool):
    """A native tool that lets a faster executor model consult a stronger advisor model mid-generation.

    The fields map 1:1 to the parameters of Anthropic's advisor tool definition. OpenRouter exposes
    the advisor as a gateway server tool that honors a subset (`model`, `max_tokens`) and ignores
    the unsupported fields; see the per-field docstrings for which provider supports each.

    Supported by:

    * Anthropic
    * OpenRouter
    """

    model: AdvisorModelName
    """The advisor model to consult, i.e. the `model` field of the provider's advisor tool definition.

    The executor/advisor pairing is validated by the provider's API, not here. The accepted namespace
    depends on the executing provider: Anthropic model IDs (e.g. `claude-opus-4-8`) on Anthropic,
    OpenRouter catalog slugs (e.g. `anthropic/claude-opus-4.8`) on OpenRouter.

    Supported by:

    * Anthropic
    * OpenRouter
    """

    max_uses: int | None = None
    """If provided, the advisor can be consulted at most this many times per request. Maps to `max_uses`.

    This is a per-request cap, not a per-run budget: a run that spans multiple requests resets the
    count each request. Enforce a conversation-wide ceiling yourself if you need one.

    Supported by:

    * Anthropic

    OpenRouter caps advisor consultations per request with a fixed gateway limit and ignores
    `max_uses`.
    """

    max_tokens: int | None = None
    """If provided, caps the advisor's output tokens (minimum 1024). Maps to `max_tokens` on Anthropic
    and `max_completion_tokens` on OpenRouter.

    When set, the Anthropic advisor result carries a `stop_reason`.

    Supported by:

    * Anthropic
    * OpenRouter
    """

    caching: Literal['5m', '1h'] | None = None
    """If provided, caches the advisor context ephemerally with the given TTL. Maps to
    `caching={'type': 'ephemeral', 'ttl': ...}`.

    Supported by:

    * Anthropic

    OpenRouter's advisor tool has no equivalent knob and ignores `caching`.
    """

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

    def __post_init__(self) -> None:
        if self.max_tokens is not None and self.max_tokens < 1024:
            raise ValueError('AdvisorTool.max_tokens must be at least 1024')


# 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


SUPPORTED_NATIVE_TOOLS = frozenset(NATIVE_TOOL_TYPES.values())
"""Set of all native tool types."""

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