# Copyright (c) Microsoft. All rights reserved.

import base64
import logging
import uuid
from asyncio import CancelledError
from collections.abc import Mapping
from functools import partial
from typing import Any

from a2a.helpers import new_task_from_user_message
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import TaskUpdater
from a2a.types import Part, TaskState
from agent_framework import (
    AgentResponseUpdate,
    AgentSession,
    Message,
    SupportsAgentRun,
)
from typing_extensions import override

from agent_framework_a2a._utils import get_uri_data

logger = logging.getLogger("agent_framework.a2a")


class A2AExecutor(AgentExecutor):
    """Execute AI agents using the A2A (Agent-to-Agent) protocol.

    The A2AExecutor bridges AI agents built with the agent_framework library and the A2A protocol,
    enabling structured agent execution with event-driven communication. It handles execution
    contexts, delegates history management to the agent's session, and converts agent
    responses into A2A protocol events.

    The executor supports executing an Agent or WorkflowAgent. It provides comprehensive
    error handling with task status updates and supports various content types including text,
    binary data, and URI-based content.

    Example:
        .. code-block:: python

            from a2a.server.request_handlers import DefaultRequestHandler
            from a2a.server.routes import create_jsonrpc_routes, create_agent_card_routes
            from a2a.server.tasks import InMemoryTaskStore
            from a2a.types import AgentCapabilities, AgentCard, AgentInterface
            from agent_framework.a2a import A2AExecutor
            from agent_framework.openai import OpenAIResponsesClient
            from starlette.applications import Starlette

            public_agent_card = AgentCard(
                name="Food Agent",
                description="A simple agent that provides food-related information.",
                version="1.0.0",
                default_input_modes=["text"],
                default_output_modes=["text"],
                capabilities=AgentCapabilities(streaming=True),
                supported_interfaces=[
                    AgentInterface(url="http://localhost:9999/", protocol_binding="JSONRPC"),
                ],
                skills=[],
            )

            # Create an agent
            agent = OpenAIResponsesClient().as_agent(
                name="Food Agent",
                instructions="A simple agent that provides food-related information.",
            )

            # Set up the A2A server with the A2AExecutor enabled for streaming
            # and passing custom keyword arguments to the agent's run method.
            request_handler = DefaultRequestHandler(
                agent_executor=A2AExecutor(agent, stream=True, run_kwargs={"client_kwargs": {"max_tokens": 500}}),
                task_store=InMemoryTaskStore(),
                agent_card=public_agent_card,
            )

            app = Starlette(
                routes=[
                    *create_agent_card_routes(public_agent_card),
                    *create_jsonrpc_routes(request_handler, "/"),
                ],
            )

    Args:
        agent: The AI agent to execute.
        stream: Whether to stream the agent response. Defaults to False.
        run_kwargs: Additional keyword arguments to pass to the agent's run method.
    """

    def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Mapping[str, Any] | None = None):
        """Initialize the A2AExecutor with the specified agent.

        Args:
            agent: The AI agent or workflow to execute.
            stream: Whether to stream the agent response. Defaults to False.
            run_kwargs: Additional keyword arguments to pass to the agent's run method.
                Cannot contain 'session' or 'stream' as these are managed by the executor.

        Raises:
            ValueError: If run_kwargs contains 'session' or 'stream'.
        """
        super().__init__()
        self._agent: SupportsAgentRun = agent
        self._stream: bool = stream
        if run_kwargs:
            if "session" in run_kwargs:
                raise ValueError("run_kwargs cannot contain 'session' as it is managed by the executor.")
            if "stream" in run_kwargs:
                raise ValueError("run_kwargs cannot contain 'stream' as it is managed by the executor.")
        self._run_kwargs: Mapping[str, Any] = run_kwargs or {}

    @override
    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Cancel agent execution for the given request context.

        Uses a TaskUpdater to send a cancellation event through the provided event queue.

        Args:
            context: The request context identifying the task to cancel.
            event_queue: The event queue to publish the cancellation event to.

        Raises:
            ValueError: If context_id is not provided in the RequestContext.
        """
        if context.context_id is None:
            raise ValueError("Context ID must be provided in the RequestContext")

        updater = TaskUpdater(
            event_queue=event_queue,
            task_id=context.task_id or "",
            context_id=context.context_id,
        )

        await updater.cancel()

    @override
    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Execute the agent with the given context and event queue.

        Orchestrates the agent execution process: sets up the agent session,
        executes the agent, processes response messages, and handles errors with appropriate task status updates.
        """
        if context.context_id is None:
            raise ValueError("Context ID must be provided in the RequestContext")
        if context.message is None:
            raise ValueError("Message must be provided in the RequestContext")

        query = context.get_user_input()
        task = context.current_task

        if not task:
            task = new_task_from_user_message(context.message)
            await event_queue.enqueue_event(task)

        updater = TaskUpdater(event_queue, task.id, context.context_id)
        await updater.submit()

        try:
            await updater.start_work()

            session = self._agent.create_session(session_id=task.context_id)

            if self._stream:
                await self._run_stream(query, session, updater)
            else:
                await self._run(query, session, updater)

            # Mark as complete
            await updater.complete()
        except CancelledError:
            await updater.update_status(state=TaskState.TASK_STATE_CANCELED)
        except Exception as e:
            logger.exception("A2AExecutor encountered an error during execution.", exc_info=e)
            await updater.update_status(
                state=TaskState.TASK_STATE_FAILED,
                message=updater.new_agent_message([Part(text=str(e))]),
            )

    async def _run_stream(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None:
        """Run the agent in streaming mode and publish updates to the task updater."""
        response_stream = self._agent.run(query, session=session, stream=True, **self._run_kwargs)
        streamed_artifact_ids: set[str] = set()
        # Generate a stable artifact ID for the entire stream so all chunks share the same ID.
        # This ensures clients can coalesce streaming tokens into a single artifact/message
        # per the A2A spec (TaskArtifactUpdateEvent with append=True on same artifactId).
        default_artifact_id = str(uuid.uuid4())
        await (
            response_stream.with_transform_hook(
                partial(
                    self.handle_events,
                    updater=updater,
                    streamed_artifact_ids=streamed_artifact_ids,
                    default_artifact_id=default_artifact_id,
                )
            )
        ).get_final_response()

    async def _run(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None:
        """Run the agent in non-streaming mode and publish messages to the task updater."""
        response = await self._agent.run(query, session=session, stream=False, **self._run_kwargs)
        response_messages = response.messages

        if not isinstance(response_messages, list):
            response_messages = [response_messages]

        for message in response_messages:
            await self.handle_events(message, updater)

    async def handle_events(
        self,
        item: Message | AgentResponseUpdate,
        updater: TaskUpdater,
        streamed_artifact_ids: set[str] | None = None,
        default_artifact_id: str | None = None,
    ) -> None:
        """Convert agent response items (Messages or Updates) to A2A protocol events.

        Processes Message or AgentResponseUpdate objects and converts them into A2A protocol format.
        Handles text, data, and URI content. USER role messages are skipped.

        Users can override this method in a subclass to implement custom transformations
        from their agent's output format to A2A protocol events.

        Args:
            item: The agent response item (Message or AgentResponseUpdate) to process.
            updater: The task updater to publish events to.
            streamed_artifact_ids: A set of artifact IDs that have already been streamed.
                Used to track which artifacts need append=True on subsequent chunks.
            default_artifact_id: A stable artifact ID to use when the item does not provide one.
                This ensures all streaming chunks for a single response share the same artifact ID,
                allowing clients to coalesce them into a single message.

        Example:
            .. code-block:: python

                class CustomA2AExecutor(A2AExecutor):
                    async def handle_events(
                        self,
                        item: Message | AgentResponseUpdate,
                        updater: TaskUpdater,
                        streamed_artifact_ids: set[str] | None = None,
                        default_artifact_id: str | None = None,
                    ) -> None:
                        # Custom logic to transform item contents
                        if item.role == "assistant" and item.contents:
                            parts = [Part(text=f"Custom: {item.contents[0].text}")]
                            await updater.update_status(
                                state=TaskState.TASK_STATE_WORKING,
                                message=updater.new_agent_message(parts=parts),
                            )
                        else:
                            await super().handle_events(item, updater)
        """
        role = getattr(item, "role", None)
        if role == "user":
            # This is a user message, we can ignore it in the context of task updates
            return

        parts: list[Part] = []
        metadata = getattr(item, "additional_properties", None)

        # AgentResponseUpdate uses 'contents', Message uses 'contents'
        contents = getattr(item, "contents", [])

        for content in contents:
            if content.type == "text" and content.text:
                parts.append(Part(text=content.text))
            elif content.type == "data" and content.uri:
                base64_str = get_uri_data(content.uri)
                parts.append(Part(raw=base64.b64decode(base64_str), media_type=content.media_type or ""))
            elif content.type == "uri" and content.uri:
                parts.append(Part(url=content.uri, media_type=content.media_type or ""))
            else:
                # Silently skip unsupported content types
                logger.warning("A2AExecutor does not yet support content type: %s. Omitted.", content.type)

        if parts:
            if isinstance(item, AgentResponseUpdate):
                # Resolve artifact ID: use item's message_id if available, otherwise fall back
                # to the stable default_artifact_id so all streaming chunks share the same ID.
                artifact_id = item.message_id or default_artifact_id
                # For streaming updates, we send TaskArtifactUpdateEvent via add_artifact
                await updater.add_artifact(
                    parts=parts,
                    artifact_id=artifact_id,
                    metadata=metadata,
                    append=(
                        True if streamed_artifact_ids is not None and artifact_id in streamed_artifact_ids else None
                    ),
                )
                if artifact_id and streamed_artifact_ids is not None:
                    streamed_artifact_ids.add(artifact_id)
            else:
                # For final messages, we send TaskStatusUpdateEvent with 'working' state
                await updater.update_status(
                    state=TaskState.TASK_STATE_WORKING,
                    message=updater.new_agent_message(parts=parts, metadata=metadata),
                )
