"""SamplingTool for use during LLM sampling requests."""

from __future__ import annotations

import inspect
from collections.abc import Callable
from typing import Any

from mcp.types import TextContent
from mcp.types import Tool as SDKTool
from pydantic import ConfigDict

from fastmcp.exceptions import AuthorizationError
from fastmcp.server.auth.authorization import AuthContext, run_auth_checks
from fastmcp.server.dependencies import get_access_token
from fastmcp.tools.base import ToolResult
from fastmcp.tools.function_parsing import ParsedFunction
from fastmcp.tools.function_tool import FunctionTool
from fastmcp.tools.tool_transform import TransformedTool
from fastmcp.utilities.types import FastMCPBaseModel


class SamplingTool(FastMCPBaseModel):
    """A tool that can be used during LLM sampling.

    SamplingTools bundle a tool's schema (name, description, parameters) with
    an executor function, enabling servers to execute agentic workflows where
    the LLM can request tool calls during sampling.

    In most cases, pass functions directly to ctx.sample():

        def search(query: str) -> str:
            '''Search the web.'''
            return web_search(query)

        result = await context.sample(
            messages="Find info about Python",
            tools=[search],  # Plain functions work directly
        )

    Create a SamplingTool explicitly when you need custom name/description:

        tool = SamplingTool.from_function(search, name="web_search")
    """

    name: str
    description: str | None = None
    parameters: dict[str, Any]
    fn: Callable[..., Any]
    sequential: bool = False

    model_config = ConfigDict(arbitrary_types_allowed=True)

    async def run(self, arguments: dict[str, Any] | None = None) -> Any:
        """Execute the tool with the given arguments.

        Args:
            arguments: Dictionary of arguments to pass to the tool function.

        Returns:
            The result of executing the tool function.
        """
        if arguments is None:
            arguments = {}

        result = self.fn(**arguments)
        if inspect.isawaitable(result):
            result = await result
        return result

    def _to_sdk_tool(self) -> SDKTool:
        """Convert to an mcp.types.Tool for SDK compatibility.

        This is used internally when passing tools to the MCP SDK's
        create_message() method.
        """
        return SDKTool(
            name=self.name,
            description=self.description,
            inputSchema=self.parameters,
        )

    @classmethod
    def from_function(
        cls,
        fn: Callable[..., Any],
        *,
        name: str | None = None,
        description: str | None = None,
        sequential: bool = False,
    ) -> SamplingTool:
        """Create a SamplingTool from a function.

        The function's signature is analyzed to generate a JSON schema for
        the tool's parameters. Type hints are used to determine parameter types.

        Args:
            fn: The function to create a tool from.
            name: Optional name override. Defaults to the function's name.
            description: Optional description override. Defaults to the function's docstring.
            sequential: If True, this tool requires sequential execution and prevents
                parallel execution of all tools in the batch. Set to True for tools
                with shared state, file writes, or other operations that cannot run
                concurrently. Defaults to False.

        Returns:
            A SamplingTool wrapping the function.

        Raises:
            ValueError: If the function is a lambda without a name override.
        """
        parsed = ParsedFunction.from_function(fn, validate=True)

        if name is None and parsed.name == "<lambda>":
            raise ValueError("You must provide a name for lambda functions")

        return cls(
            name=name or parsed.name,
            description=description if description is not None else parsed.description,
            parameters=parsed.input_schema,
            fn=parsed.fn,
            sequential=sequential,
        )

    @classmethod
    def from_callable_tool(
        cls,
        tool: FunctionTool | TransformedTool,
        *,
        name: str | None = None,
        description: str | None = None,
    ) -> SamplingTool:
        """Create a SamplingTool from a FunctionTool or TransformedTool.

        Reuses existing server tools in sampling contexts. For TransformedTool,
        the tool's .run() method is used to ensure proper argument transformation,
        and the ToolResult is automatically unwrapped.

        Args:
            tool: A FunctionTool or TransformedTool to convert.
            name: Optional name override. Defaults to tool.name.
            description: Optional description override. Defaults to tool.description.

        Raises:
            TypeError: If the tool is not a FunctionTool or TransformedTool.
        """
        # Validate that the tool is a supported type
        if not isinstance(tool, (FunctionTool, TransformedTool)):
            raise TypeError(
                f"Expected FunctionTool or TransformedTool, got {type(tool).__name__}. "
                "Only callable tools can be converted to SamplingTools."
            )

        # Both FunctionTool and TransformedTool need .run() to ensure proper
        # result processing (serializers, output_schema, wrap-result flags)
        async def wrapper(**kwargs: Any) -> Any:
            # Enforce per-tool auth checks, mirroring what the server
            # dispatcher does for direct tool calls.  Without this, an
            # auth-protected tool wrapped as a SamplingTool could be
            # invoked by the LLM during sampling without authorization.
            if tool.auth is not None:
                # Late import to avoid circular import with context.py
                from fastmcp.server.context import _current_transport

                is_stdio = _current_transport.get() == "stdio"
                if not is_stdio:
                    token = get_access_token()
                    ctx = AuthContext(token=token, component=tool)
                    if not await run_auth_checks(tool.auth, ctx):
                        raise AuthorizationError(
                            f"Authorization failed for tool '{tool.name}': "
                            "insufficient permissions"
                        )

            result = await tool.run(kwargs)
            # Unwrap ToolResult - extract the actual value
            if isinstance(result, ToolResult):
                # If there's structured_content, use that
                if result.structured_content is not None:
                    # Check tool's schema - this is the source of truth
                    if tool.output_schema and tool.output_schema.get(
                        "x-fastmcp-wrap-result"
                    ):
                        # Tool wraps results: {"result": value} -> value
                        return result.structured_content.get("result")
                    else:
                        # No wrapping: use structured_content directly
                        return result.structured_content
                # Otherwise, extract from text content
                if result.content and len(result.content) > 0:
                    first_content = result.content[0]
                    if isinstance(first_content, TextContent):
                        return first_content.text
            return result

        fn = wrapper

        # Extract the callable function, name, description, and parameters
        return cls(
            name=name or tool.name,
            description=description or tool.description,
            parameters=tool.parameters,
            fn=fn,
        )
