import asyncio
import builtins
import json
from abc import ABC
from typing import Any, Dict, Generic, Sequence, Type, TypeVar

from agentdhal_core import CancellationToken
from agentdhal_core.tools import BaseTool
from agentdhal_core.utils import schema_to_pydantic_model
from pydantic import BaseModel
from pydantic.networks import AnyUrl

from mcp import ClientSession, Tool
from mcp.types import AudioContent, ContentBlock, EmbeddedResource, ImageContent, ResourceLink, TextContent

from ._config import McpServerParams
from ._session import create_mcp_server_session

TServerParams = TypeVar("TServerParams", bound=McpServerParams)


class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
    """
    Base adapter class for MCP tools to make them compatible with AutoGen.

    Args:
        server_params (TServerParams): Parameters for the MCP server connection.
        tool (Tool): The MCP tool to wrap.
    """

    component_type = "tool"

    def __init__(self, server_params: TServerParams, tool: Tool, session: ClientSession | None = None) -> None:
        self._tool = tool
        self._server_params = server_params
        self._session = session

        # Extract name and description
        name = tool.name
        description = tool.description or ""

        # Create the input model from the tool's schema
        input_model = schema_to_pydantic_model(tool.inputSchema)

        # Use Any as return type since MCP tool returns can vary
        return_type: Type[Any] = object

        super().__init__(input_model, return_type, name, description)

    async def run(self, args: BaseModel, cancellation_token: CancellationToken) -> Any:
        """
        Run the MCP tool with the provided arguments.

        Args:
            args (BaseModel): The arguments to pass to the tool.
            cancellation_token (CancellationToken): Token to signal cancellation.

        Returns:
            Any: The result of the tool execution.

        Raises:
            Exception: If the operation is cancelled or the tool execution fails.
        """
        # Convert the input model to a dictionary
        # Exclude unset values to avoid sending them to the MCP servers which may cause errors
        # for many servers.
        kwargs = args.model_dump(exclude_unset=True)

        if self._session is not None:
            # If a session is provided, use it directly.
            session = self._session
            return await self._run(args=kwargs, cancellation_token=cancellation_token, session=session)

        async with create_mcp_server_session(self._server_params) as session:
            await session.initialize()
            return await self._run(args=kwargs, cancellation_token=cancellation_token, session=session)

    def _normalize_payload_to_content_list(self, payload: Sequence[ContentBlock]) -> list[ContentBlock]:
        """
        Normalizes a raw tool output payload into a list of content items.
        - If payload is already a sequence of ContentBlock items, it's converted to a list and returned.
        - If payload is a single ContentBlock item, it's wrapped in a list.
        - If payload is a string, it's wrapped in [TextContent(text=payload)].
        - Otherwise, the payload is stringified and wrapped in [TextContent(text=str(payload))].
        """
        if isinstance(payload, Sequence) and all(
            isinstance(item, (TextContent, ImageContent, EmbeddedResource, AudioContent, ResourceLink))
            for item in payload
        ):
            return list(payload)
        elif isinstance(payload, (TextContent, ImageContent, EmbeddedResource, AudioContent, ResourceLink)):
            return [payload]
        elif isinstance(payload, str):
            return [TextContent(text=payload, type="text")]
        else:
            return [TextContent(text=str(payload), type="text")]

    async def _run(self, args: Dict[str, Any], cancellation_token: CancellationToken, session: ClientSession) -> Any:
        exceptions_to_catch: tuple[Type[BaseException], ...]
        if hasattr(builtins, "ExceptionGroup"):
            exceptions_to_catch = (asyncio.CancelledError, builtins.ExceptionGroup)
        else:
            exceptions_to_catch = (asyncio.CancelledError,)

        try:
            if cancellation_token.is_cancelled():
                raise asyncio.CancelledError("Operation cancelled")

            result_future = asyncio.ensure_future(session.call_tool(name=self._tool.name, arguments=args))
            cancellation_token.link_future(result_future)
            result = await result_future

            normalized_content_list = self._normalize_payload_to_content_list(result.content)

            if result.isError:
                serialized_error_message = self.return_value_as_string(normalized_content_list)
                raise Exception(serialized_error_message)
            return normalized_content_list

        except exceptions_to_catch:
            # Re-raise these specific exception types directly.
            raise

    @classmethod
    async def from_server_params(cls, server_params: TServerParams, tool_name: str) -> "McpToolAdapter[TServerParams]":
        """
        Create an instance of McpToolAdapter from server parameters and tool name.

        Args:
            server_params (TServerParams): Parameters for the MCP server connection.
            tool_name (str): The name of the tool to wrap.

        Returns:
            McpToolAdapter[TServerParams]: An instance of McpToolAdapter.

        Raises:
            ValueError: If the tool with the specified name is not found.
        """
        async with create_mcp_server_session(server_params) as session:
            await session.initialize()

            tools_response = await session.list_tools()
            matching_tool = next((t for t in tools_response.tools if t.name == tool_name), None)

            if matching_tool is None:
                raise ValueError(
                    f"Tool '{tool_name}' not found, available tools: {', '.join([t.name for t in tools_response.tools])}"
                )

        return cls(server_params=server_params, tool=matching_tool)

    def return_value_as_string(self, value: list[Any]) -> str:
        """Return a string representation of the result."""

        def serialize_item(item: Any) -> dict[str, Any]:
            if isinstance(item, (TextContent, ImageContent, AudioContent)):
                dumped = item.model_dump()
                # Remove the 'meta' field if it exists and is None (for backward compatibility)
                if dumped.get("meta") is None:
                    dumped.pop("meta", None)
                return dumped
            elif isinstance(item, EmbeddedResource):
                type = item.type
                resource = {}
                for key, val in item.resource.model_dump().items():
                    # Skip 'meta' field if it's None (for backward compatibility)
                    if key == "meta" and val is None:
                        continue
                    if isinstance(val, AnyUrl):
                        resource[key] = str(val)
                    else:
                        resource[key] = val
                dumped_annotations = item.annotations.model_dump() if item.annotations else None
                # Remove 'meta' from annotations if it exists and is None
                if dumped_annotations and dumped_annotations.get("meta") is None:
                    dumped_annotations.pop("meta", None)
                return {"type": type, "resource": resource, "annotations": dumped_annotations}
            elif isinstance(item, ResourceLink):
                dumped = item.model_dump()
                # Remove the 'meta' field if it exists and is None (for backward compatibility)
                if dumped.get("meta") is None:
                    dumped.pop("meta", None)
                # Convert AnyUrl to string for JSON serialization
                if "uri" in dumped and isinstance(dumped["uri"], AnyUrl):
                    dumped["uri"] = str(dumped["uri"])
                return dumped
            else:
                return {}

        return json.dumps([serialize_item(item) for item in value])
