from __future__ import annotations

from collections.abc import Awaitable, Callable
from typing import Any, Concatenate, ParamSpec, TypeVar, overload

from pydantic.json_schema import GenerateJsonSchema
from pydantic_ai import Agent
from pydantic_ai.tools import (
    ArgsValidatorFunc,
    DocstringFormat,
    GenerateToolJsonSchema,
    RunContext,
    ToolPrepareFunc,
)

AgentDepsT = TypeVar("AgentDepsT")
ToolParamsT = ParamSpec("ToolParamsT")
ToolReturnT = TypeVar("ToolReturnT")


@overload
def typed_tool(
    agent: Agent[AgentDepsT, Any],
    func: Callable[
        Concatenate[RunContext[AgentDepsT], ToolParamsT],
        ToolReturnT | Awaitable[ToolReturnT],
    ],
    /,
    *,
    name: str | None = None,
    description: str | None = None,
    retries: int | None = None,
    prepare: ToolPrepareFunc[AgentDepsT] | None = None,
    args_validator: ArgsValidatorFunc[AgentDepsT, ToolParamsT] | None = None,
    docstring_format: DocstringFormat = "auto",
    require_parameter_descriptions: bool = False,
    schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
    strict: bool | None = None,
    sequential: bool = False,
    requires_approval: bool = False,
    metadata: dict[str, Any] | None = None,
    timeout: float | None = None,
) -> Callable[
    Concatenate[RunContext[AgentDepsT], ToolParamsT],
    ToolReturnT | Awaitable[ToolReturnT],
]: ...


@overload
def typed_tool(
    agent: Agent[AgentDepsT, Any],
    func: None = None,
    /,
    *,
    name: str | None = None,
    description: str | None = None,
    retries: int | None = None,
    prepare: ToolPrepareFunc[AgentDepsT] | None = None,
    args_validator: ArgsValidatorFunc[AgentDepsT, ToolParamsT] | None = None,
    docstring_format: DocstringFormat = "auto",
    require_parameter_descriptions: bool = False,
    schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
    strict: bool | None = None,
    sequential: bool = False,
    requires_approval: bool = False,
    metadata: dict[str, Any] | None = None,
    timeout: float | None = None,
) -> Callable[
    [
        Callable[
            Concatenate[RunContext[AgentDepsT], ToolParamsT],
            ToolReturnT | Awaitable[ToolReturnT],
        ]
    ],
    Callable[
        Concatenate[RunContext[AgentDepsT], ToolParamsT],
        ToolReturnT | Awaitable[ToolReturnT],
    ],
]: ...


def typed_tool(
    agent: Agent[AgentDepsT, Any],
    func: Callable[..., Any] | None = None,
    /,
    *,
    name: str | None = None,
    description: str | None = None,
    retries: int | None = None,
    prepare: ToolPrepareFunc[AgentDepsT] | None = None,
    args_validator: ArgsValidatorFunc[AgentDepsT, ToolParamsT] | None = None,
    docstring_format: DocstringFormat = "auto",
    require_parameter_descriptions: bool = False,
    schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
    strict: bool | None = None,
    sequential: bool = False,
    requires_approval: bool = False,
    metadata: dict[str, Any] | None = None,
    timeout: float | None = None,
) -> Callable[[Callable[..., Any]], Callable[..., Any]] | Callable[..., Any]:
    """Register `agent.tool(...)` without losing the decorated function type."""

    def decorator(func_: Callable[..., Any]) -> Callable[..., Any]:
        agent.tool(
            name=name,
            description=description,
            retries=retries,
            prepare=prepare,
            args_validator=args_validator,
            docstring_format=docstring_format,
            require_parameter_descriptions=require_parameter_descriptions,
            schema_generator=schema_generator,
            strict=strict,
            sequential=sequential,
            requires_approval=requires_approval,
            metadata=metadata,
            timeout=timeout,
        )(func_)
        return func_

    return decorator if func is None else decorator(func)


@overload
def typed_tool_plain(
    agent: Agent[AgentDepsT, Any],
    func: Callable[ToolParamsT, ToolReturnT | Awaitable[ToolReturnT]],
    /,
    *,
    name: str | None = None,
    description: str | None = None,
    retries: int | None = None,
    prepare: ToolPrepareFunc[AgentDepsT] | None = None,
    args_validator: ArgsValidatorFunc[AgentDepsT, ToolParamsT] | None = None,
    docstring_format: DocstringFormat = "auto",
    require_parameter_descriptions: bool = False,
    schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
    strict: bool | None = None,
    sequential: bool = False,
    requires_approval: bool = False,
    metadata: dict[str, Any] | None = None,
    timeout: float | None = None,
) -> Callable[ToolParamsT, ToolReturnT | Awaitable[ToolReturnT]]: ...


@overload
def typed_tool_plain(
    agent: Agent[AgentDepsT, Any],
    func: None = None,
    /,
    *,
    name: str | None = None,
    description: str | None = None,
    retries: int | None = None,
    prepare: ToolPrepareFunc[AgentDepsT] | None = None,
    args_validator: ArgsValidatorFunc[AgentDepsT, ToolParamsT] | None = None,
    docstring_format: DocstringFormat = "auto",
    require_parameter_descriptions: bool = False,
    schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
    strict: bool | None = None,
    sequential: bool = False,
    requires_approval: bool = False,
    metadata: dict[str, Any] | None = None,
    timeout: float | None = None,
) -> Callable[
    [Callable[ToolParamsT, ToolReturnT | Awaitable[ToolReturnT]]],
    Callable[ToolParamsT, ToolReturnT | Awaitable[ToolReturnT]],
]: ...


def typed_tool_plain(
    agent: Agent[AgentDepsT, Any],
    func: Callable[..., Any] | None = None,
    /,
    *,
    name: str | None = None,
    description: str | None = None,
    retries: int | None = None,
    prepare: ToolPrepareFunc[AgentDepsT] | None = None,
    args_validator: ArgsValidatorFunc[AgentDepsT, ToolParamsT] | None = None,
    docstring_format: DocstringFormat = "auto",
    require_parameter_descriptions: bool = False,
    schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
    strict: bool | None = None,
    sequential: bool = False,
    requires_approval: bool = False,
    metadata: dict[str, Any] | None = None,
    timeout: float | None = None,
) -> Callable[[Callable[..., Any]], Callable[..., Any]] | Callable[..., Any]:
    """Register `agent.tool_plain(...)` without losing the decorated function type."""

    def decorator(func_: Callable[..., Any]) -> Callable[..., Any]:
        agent.tool_plain(
            name=name,
            description=description,
            retries=retries,
            prepare=prepare,
            args_validator=args_validator,
            docstring_format=docstring_format,
            require_parameter_descriptions=require_parameter_descriptions,
            schema_generator=schema_generator,
            strict=strict,
            sequential=sequential,
            requires_approval=requires_approval,
            metadata=metadata,
            timeout=timeout,
        )(func_)
        return func_

    return decorator if func is None else decorator(func)
