"""
MCP Tool Wrappers for Pydantic AI.

This module provides functions to convert MCP tool functions and resources
into a format compatible with the Pydantic AI library.

Context Propagation:
    MCP tools need access to AgentContext (user_id, session_id, etc.) via
    get_current_context(). This module captures the context at tool creation
    time and ensures it's set when tools execute, working around contextvar
    propagation issues in async tool execution.
"""

from contextvars import copy_context
from typing import Any, Callable

from loguru import logger
from pydantic_ai.tools import Tool


def create_pydantic_tool(func: Callable[..., Any]) -> Tool:
    """
    Create a Pydantic AI Tool from a given function.

    This uses the Tool constructor, which inspects the
    function's signature and docstring to create the tool schema.

    Args:
        func: The function to wrap as a tool.

    Returns:
        A Pydantic AI Tool instance.
    """
    logger.debug(f"Creating Pydantic tool from function: {func.__name__}")
    return Tool(func)


def create_mcp_tool_wrapper(
    tool_name: str,
    mcp_tool: Any,
    user_id: str | None = None,
    description_suffix: str | None = None,
    agent_context: Any = None,
) -> Tool:
    """
    Create a Pydantic AI Tool from a FastMCP FunctionTool.

    FastMCP tools are FunctionTool objects that wrap the actual async function.
    We pass the function directly to Pydantic AI's Tool class, which will
    inspect its signature properly. User ID injection happens in the wrapper.

    Args:
        tool_name: Name of the MCP tool
        mcp_tool: The FastMCP FunctionTool object
        user_id: Optional user_id to inject into tool calls
        description_suffix: Optional text to append to the tool's docstring.
            Used to add schema-specific context (e.g., default table for search_rem).
        agent_context: Optional AgentContext to propagate to tool execution.
            Ensures get_current_context() works in MCP tools.

    Returns:
        A Pydantic AI Tool instance
    """
    # Extract the actual function from FastMCP FunctionTool
    tool_func = mcp_tool.fn

    # Check if function accepts user_id parameter
    import inspect
    sig = inspect.signature(tool_func)
    has_user_id = "user_id" in sig.parameters

    # Build the docstring with optional suffix
    base_doc = tool_func.__doc__ or ""
    final_doc = base_doc + description_suffix if description_suffix else base_doc

    # Capture context for propagation (closure captures reference)
    _captured_context = agent_context

    # Always create a wrapper to ensure context propagation
    async def wrapped_tool(**kwargs) -> Any:
        """Wrapper that propagates context and optionally injects user_id."""
        from ..context import get_current_context, set_current_context

        # Set context for tool execution if provided and not already set
        previous_context = get_current_context()
        if _captured_context is not None and previous_context is None:
            set_current_context(_captured_context)
            logger.debug(f"Set context for tool {tool_name}: user_id={_captured_context.user_id}")

        try:
            # Inject user_id if function accepts it and it's not provided
            if has_user_id and "user_id" not in kwargs and user_id:
                kwargs["user_id"] = user_id
                logger.debug(f"Injecting user_id={user_id} into tool {tool_name}")

            # Filter kwargs to only include parameters that the function accepts
            valid_params = set(sig.parameters.keys())
            filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}

            return await tool_func(**filtered_kwargs)
        finally:
            # Restore previous context
            if _captured_context is not None and previous_context is None:
                set_current_context(previous_context)

    # Copy signature from original function for Pydantic AI inspection
    wrapped_tool.__name__ = tool_name
    wrapped_tool.__doc__ = final_doc
    wrapped_tool.__annotations__ = tool_func.__annotations__
    wrapped_tool.__signature__ = sig  # Important: preserve full signature

    logger.debug(f"Creating MCP tool wrapper: {tool_name} (context={'set' if agent_context else 'None'})")
    return Tool(wrapped_tool)


def create_resource_tool(uri: str, usage: str = "", mcp_server: Any = None, agent_context: Any = None) -> Tool:
    """
    Build a Tool instance from an MCP resource URI.

    Creates a tool that fetches the resource content when called.
    Resources declared in agent YAML become callable tools - this eliminates
    the artificial MCP distinction between tools and resources.

    Supports both:
    - Concrete URIs: "rem://agents" -> tool with no parameters
    - Template URIs: "patient-profile://field/{field_key}" -> tool with field_key parameter

    Args:
        uri: The resource URI (concrete or template with {variable} placeholders).
        usage: The description of what this resource provides.
        mcp_server: Optional FastMCP server instance to resolve resources from.
            If provided, resources are resolved from this server's registry.
            If not provided, falls back to REM's built-in load_resource().
        agent_context: Optional AgentContext to propagate to tool execution.
            Ensures get_current_context() works in MCP resources.

    Returns:
        A Pydantic AI Tool instance that fetches the resource.

    Example:
        # Concrete URI -> no-param tool
        tool = create_resource_tool("rem://agents", "List all agent schemas")

        # Template URI -> parameterized tool
        tool = create_resource_tool("patient-profile://field/{field_key}", "Get field definition", mcp_server=mcp)
        # Agent calls: get_patient_profile_field(field_key="safety.suicidality")
    """
    import json
    import re

    # Extract template variables from URI (e.g., {field_key}, {domain_name})
    template_vars = re.findall(r'\{([^}]+)\}', uri)

    # Parse URI to create function name (strip template vars for cleaner name)
    clean_uri = re.sub(r'\{[^}]+\}', '', uri)
    parts = clean_uri.replace("://", "_").replace("-", "_").replace("/", "_").replace(".", "_")
    parts = re.sub(r'_+', '_', parts).strip('_')  # Clean up multiple underscores
    func_name = f"get_{parts}"

    # For parameterized URIs, append _by_{params} to avoid naming conflicts
    # e.g., rem://agents/{name} -> get_rem_agents_by_name (distinct from get_rem_agents)
    if template_vars:
        param_suffix = "_by_" + "_".join(template_vars)
        func_name = f"{func_name}{param_suffix}"

    # Build description including parameter info
    description = usage or f"Fetch {uri} resource"
    if template_vars:
        param_desc = ", ".join(template_vars)
        description = f"{description}\n\nParameters: {param_desc}"

    # Capture references at tool creation time (for closure)
    # This ensures the correct server/context is used even if called later
    _captured_mcp_server = mcp_server
    _captured_uri = uri  # Also capture URI for consistent logging
    _captured_context = agent_context  # Capture context for propagation

    if template_vars:
        # Template URI -> create parameterized tool
        async def wrapper(**kwargs: Any) -> str:
            """Fetch MCP resource with substituted parameters."""
            import asyncio
            import inspect
            from ..context import get_current_context, set_current_context

            # Set context for resource execution if provided and not already set
            previous_context = get_current_context()
            if _captured_context is not None and previous_context is None:
                set_current_context(_captured_context)
                logger.debug(f"Set context for resource {_captured_uri}: user_id={_captured_context.user_id}")

            try:
                logger.debug(f"Resource tool invoked: uri={_captured_uri}, kwargs={kwargs}, mcp_server={'set' if _captured_mcp_server else 'None'}")

                # Try to resolve from MCP server's resource templates first
                if _captured_mcp_server is not None:
                    try:
                        # Get resource templates from MCP server
                        templates = await _captured_mcp_server.get_resource_templates()
                        logger.debug(f"MCP server templates: {list(templates.keys())}")
                        if _captured_uri in templates:
                            template = templates[_captured_uri]
                            logger.debug(f"Found template for {_captured_uri}, calling fn with kwargs={kwargs}")
                            # Call the template's underlying function directly
                            # The fn expects the template variables as kwargs
                            fn_result = template.fn(**kwargs)
                            # Handle both sync and async functions
                            if inspect.iscoroutine(fn_result):
                                fn_result = await fn_result
                            if isinstance(fn_result, str):
                                return fn_result
                            return json.dumps(fn_result, indent=2)
                        else:
                            logger.warning(f"Template {_captured_uri} not found in MCP server templates: {list(templates.keys())}")
                    except Exception as e:
                        logger.warning(f"Failed to resolve resource {_captured_uri} from MCP server: {e}", exc_info=True)
                else:
                    logger.warning(f"No MCP server provided for resource tool {_captured_uri}, using fallback")

                # Fallback: substitute template variables and use load_resource
                resolved_uri = _captured_uri
                for var in template_vars:
                    if var in kwargs:
                        resolved_uri = resolved_uri.replace(f"{{{var}}}", str(kwargs[var]))
                    else:
                        return json.dumps({"error": f"Missing required parameter: {var}"})

                logger.debug(f"Using fallback load_resource for resolved URI: {resolved_uri}")
                from rem.api.mcp_router.resources import load_resource
                result = await load_resource(resolved_uri)
                if isinstance(result, str):
                    return result
                return json.dumps(result, indent=2)
            finally:
                # Restore previous context
                if _captured_context is not None and previous_context is None:
                    set_current_context(previous_context)

        # Build parameter annotations for Pydantic AI
        wrapper.__name__ = func_name
        wrapper.__doc__ = description
        # Add type hints for parameters
        wrapper.__annotations__ = {var: str for var in template_vars}
        wrapper.__annotations__['return'] = str

        logger.info(f"Built parameterized resource tool: {func_name} (uri: {uri}, params: {template_vars}, mcp_server={'provided' if mcp_server else 'None'})")
    else:
        # Concrete URI -> no-param tool
        async def wrapper(**kwargs: Any) -> str:
            """Fetch MCP resource and return contents."""
            import asyncio
            import inspect
            from ..context import get_current_context, set_current_context

            if kwargs:
                logger.warning(f"Resource tool {func_name} called with unexpected kwargs: {list(kwargs.keys())}")

            # Set context for resource execution if provided and not already set
            previous_context = get_current_context()
            if _captured_context is not None and previous_context is None:
                set_current_context(_captured_context)
                logger.debug(f"Set context for resource {_captured_uri}: user_id={_captured_context.user_id}")

            try:
                logger.debug(f"Concrete resource tool invoked: uri={_captured_uri}, mcp_server={'set' if _captured_mcp_server else 'None'}")

                # Try to resolve from MCP server's resources first
                if _captured_mcp_server is not None:
                    try:
                        resources = await _captured_mcp_server.get_resources()
                        logger.debug(f"MCP server resources: {list(resources.keys())}")
                        if _captured_uri in resources:
                            resource = resources[_captured_uri]
                            logger.debug(f"Found resource for {_captured_uri}")
                            # Call the resource's underlying function
                            fn_result = resource.fn()
                            if inspect.iscoroutine(fn_result):
                                fn_result = await fn_result
                            if isinstance(fn_result, str):
                                return fn_result
                            return json.dumps(fn_result, indent=2)
                        else:
                            logger.warning(f"Resource {_captured_uri} not found in MCP server resources: {list(resources.keys())}")
                    except Exception as e:
                        logger.warning(f"Failed to resolve resource {_captured_uri} from MCP server: {e}", exc_info=True)
                else:
                    logger.warning(f"No MCP server provided for resource tool {_captured_uri}, using fallback")

                # Fallback to load_resource
                logger.debug(f"Using fallback load_resource for URI: {_captured_uri}")
                from rem.api.mcp_router.resources import load_resource
                result = await load_resource(_captured_uri)
                if isinstance(result, str):
                    return result
                return json.dumps(result, indent=2)
            finally:
                # Restore previous context
                if _captured_context is not None and previous_context is None:
                    set_current_context(previous_context)

        wrapper.__name__ = func_name
        wrapper.__doc__ = description

        logger.info(f"Built resource tool: {func_name} (uri: {uri}, mcp_server={'provided' if mcp_server else 'None'})")

    return Tool(wrapper)
