# adapted from https://github.com/google/a2a-python/blob/main/examples/helloworld/test_client.py

import re
from collections.abc import Callable, Coroutine
from contextlib import suppress
from typing import TYPE_CHECKING, Any
from uuid import uuid4

from any_agent.utils.asyncio_sync import run_async_in_sync

if TYPE_CHECKING:
    from a2a.types import AgentCard

a2a_tool_available = False
with suppress(ImportError):
    import httpx
    from a2a.client import A2ACardResolver, A2AClient
    from a2a.types import (
        Message,
        MessageSendParams,
        Part,
        Role,
        SendMessageRequest,
        TextPart,
    )

    a2a_tool_available = True


async def a2a_tool_async(
    url: str, toolname: str | None = None, http_kwargs: dict[str, Any] | None = None
) -> Callable[[str], Coroutine[Any, Any, str]]:
    """Perform a query using A2A to another agent.

    Args:
        url (str): The url in which the A2A agent is located.
        toolname (str): The name for the created tool. Defaults to `call_{agent name in card}`.
            Leading and trailing whitespace are removed. Whitespace in the middle is replaced by `_`.
        http_kwargs (dict): Additional kwargs to pass to the httpx client.

    Returns:
        An async `Callable` that takes a query and returns the agent response.

    """
    if not a2a_tool_available:
        msg = "You need to `pip install 'any-agent[a2a]'` to use this tool"
        raise ImportError(msg)

    async with httpx.AsyncClient(follow_redirects=True) as resolver_client:
        a2a_agent_card: AgentCard = await (
            A2ACardResolver(httpx_client=resolver_client, base_url=url)
        ).get_agent_card()

    async def _send_query(query: str, task_id: str = str(uuid4())) -> str:
        async with httpx.AsyncClient(follow_redirects=True) as query_client:
            client = A2AClient(httpx_client=query_client, agent_card=a2a_agent_card)
            send_message_payload = SendMessageRequest(
                id=str(uuid4()),
                params=MessageSendParams(
                    message=Message(
                        role=Role.user,
                        parts=[Part(root=TextPart(text=query))],
                        # the id is not currently tracked
                        messageId=str(uuid4().hex),
                        taskId=task_id,
                    )
                ),
            )
            # TODO check how to capture exceptions and pass them on to the enclosing framework
            response = await client.send_message(
                send_message_payload, http_kwargs=http_kwargs
            )

            if not response.root:
                msg = (
                    "The A2A agent did not return a root. Are you using an A2A agent not managed by any-agent? "
                    "Please file an issue at https://github.com/mozilla-ai/any-agent/issues so we can help."
                )
                raise ValueError(msg)

            if hasattr(response.root, "error"):
                response_str = ""
                response_str += f"Error: {response.root.error.message}\n\n"
                response_str += f"Code: {response.root.error.code}\n\n"
                response_str += f"Data: {response.root.error.data}\n\n"
            elif hasattr(response.root, "result"):
                response_str = ""
                response_str += f"Status: {response.root.result.status.state}\n\n"
                response_str += f"""Message: {" ".join([part.root.text for part in response.root.result.status.message.parts if part.root.kind == "text"])}\n\n"""
                response_str += (
                    f"TaskId: {response.root.result.status.message.taskId}\n\n"
                )
                response_str += (
                    f"Timestamp: {response.root.result.status.timestamp}\n\n"
                )
            else:
                msg = (
                    "The A2A agent did not return a error or a result. Are you using an A2A agent not managed by any-agent? "
                    "Please file an issue at https://github.com/mozilla-ai/any-agent/issues so we can help."
                )
                raise ValueError(msg)

            return response_str

    new_name = toolname or a2a_agent_card.name
    new_name = re.sub(r"\s+", "_", new_name.strip())
    _send_query.__name__ = f"call_{new_name}"
    _send_query.__doc__ = f"""{a2a_agent_card.description}
        Send a query to the A2A hosted agent named {a2a_agent_card.name}.

        Agent description: {a2a_agent_card.description}

        Args:
            query (str): The query to perform.
            task_id (str): The task id to use for the conversation. Defaults to a new uuid. If you want to continue the conversation, you should provide the same task id that you received in a previous response.

        Returns:
            The result from the A2A agent, encoded as a json string.
    """
    return _send_query


def a2a_tool(
    url: str, toolname: str | None = None, http_kwargs: dict[str, Any] | None = None
) -> Callable[[str], str]:
    """Perform a query using A2A to another agent (synchronous version).

    Args:
        url (str): The url in which the A2A agent is located.
        toolname (str): The name for the created tool. Defaults to `call_{agent name in card}`.
            Leading and trailing whitespace are removed. Whitespace in the middle is replaced by `_`.
        http_kwargs (dict): Additional kwargs to pass to the httpx client.

    Returns:
        A sync `Callable` that takes a query and returns the agent response.

    """
    if not a2a_tool_available:
        msg = "You need to `pip install 'any-agent[a2a]'` to use this tool"
        raise ImportError(msg)

    # Fetch the async tool upfront to get proper name and documentation (otherwise the tool doesn't have the right name and documentation)
    async_tool = run_async_in_sync(a2a_tool_async(url, toolname, http_kwargs))

    def sync_wrapper(query: str) -> Any:
        """Execute the A2A tool query synchronously."""
        return run_async_in_sync(async_tool(query))

    # Copy essential metadata from the async tool
    sync_wrapper.__name__ = async_tool.__name__
    sync_wrapper.__doc__ = async_tool.__doc__

    return sync_wrapper
