"""A2A `AgentExecutor` that runs a compiled LangGraph workflow."""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

from a2a.server.agent_execution import AgentExecutor
from a2a.server.agent_execution.context import RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import TaskUpdater
from a2a.types import TaskState
from a2a.utils import new_agent_text_message
from langchain_core.messages import HumanMessage

from a2a_langgraph.messages import langgraph_result_to_text


def default_input_mapper(context: RequestContext) -> dict[str, Any]:
    """Map an A2A request to a LangGraph state dict with a `messages` key."""
    text = context.get_user_input()
    return {"messages": [HumanMessage(content=text)]}


class LangGraphAgentExecutor(AgentExecutor):
    """Execute a compiled LangGraph and reply with a single agent text message."""

    def __init__(
        self,
        graph: Any,
        *,
        input_mapper: Callable[[RequestContext], dict[str, Any]] | None = None,
        output_mapper: Callable[[Any], str] | None = None,
        invoke_config_builder: Callable[[RequestContext], dict[str, Any]] | None = None,
    ) -> None:
        self.graph = graph
        self.input_mapper = input_mapper or default_input_mapper
        self.output_mapper = output_mapper or langgraph_result_to_text
        self.invoke_config_builder = invoke_config_builder

    def _build_invoke_config(self, context: RequestContext) -> dict[str, Any] | None:
        if self.invoke_config_builder is not None:
            cfg = self.invoke_config_builder(context)
            return cfg if cfg else None
        if context.context_id:
            return {"configurable": {"thread_id": context.context_id}}
        return None

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        state = self.input_mapper(context)
        cfg = self._build_invoke_config(context)
        if cfg is not None:
            result = await self.graph.ainvoke(state, config=cfg)
        else:
            result = await self.graph.ainvoke(state)
        text_out = self.output_mapper(result)
        await event_queue.enqueue_event(
            new_agent_text_message(
                text_out,
                context_id=context.context_id,
                task_id=context.task_id,
            )
        )

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        task_id = context.task_id
        context_id = context.context_id
        if not task_id or not context_id:
            return
        updater = TaskUpdater(event_queue, task_id, context_id)
        await updater.update_status(TaskState.canceled, final=True)
