"""AI Agent node."""

from __future__ import annotations
import asyncio
import logging
import random
import re
from collections.abc import Mapping
from typing import Any, ClassVar, cast
from langchain.agents import create_agent
from langchain.agents.structured_output import ProviderStrategy
from langchain.chat_models import init_chat_model
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.tools import BaseTool, StructuredTool
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.graph import StateGraph
from pydantic import BaseModel, ConfigDict, Field, field_serializer
from pydantic.json_schema import SkipJsonSchema
from agentensor.tensor import TextTensor
from orcheo.graph.state import State
from orcheo.nodes.agent_tools.context import (
    get_active_tool_config,
    get_active_tool_progress_callback,
    tool_execution_context,
)
from orcheo.nodes.agent_tools.registry import tool_registry
from orcheo.nodes.base import AINode, TaskNode
from orcheo.nodes.registry import NodeMetadata, registry
from orcheo.nodes.storage import get_graph_store as _get_graph_store_fn


logger = logging.getLogger(__name__)


async def _run_tool_graph(
    compiled_graph: Runnable,
    payload: dict[str, Any],
) -> Any:
    """Execute a compiled graph, streaming updates when configured."""
    config = get_active_tool_config()
    progress_callback = get_active_tool_progress_callback()

    if progress_callback is None or config is None:
        if config is None:
            return await compiled_graph.ainvoke(payload)
        return await compiled_graph.ainvoke(payload, config=config)

    last_values: Any | None = None
    async for event in compiled_graph.astream(
        payload,
        config=config,  # type: ignore[arg-type]
        stream_mode=["updates", "values"],
    ):
        if isinstance(event, tuple) and len(event) == 2:
            mode, data = event
        else:  # pragma: no cover - defensive fallback
            mode, data = "updates", event
        if mode == "updates":
            await progress_callback(data)
        elif mode == "values":  # pragma: no branch
            last_values = data

    if last_values is not None:
        return last_values

    msg = "Tool graph streaming did not produce final values."
    raise RuntimeError(msg)


def _create_workflow_tool_func(
    compiled_graph: Runnable,
    name: str,
    description: str,
    args_schema: type[BaseModel] | None,
) -> StructuredTool:
    """Create a StructuredTool from a compiled workflow graph.

    This factory function properly binds the compiled_graph to avoid
    closure issues in loops.

    Args:
        compiled_graph: Compiled LangGraph runnable
        name: Tool name
        description: Tool description
        args_schema: Optional Pydantic model for tool arguments

    Returns:
        StructuredTool instance wrapping the workflow
    """

    async def workflow_coroutine(**kwargs: Any) -> Any:
        """Execute the workflow graph asynchronously."""
        payload = {"inputs": kwargs, "results": {}, "messages": []}
        return await _run_tool_graph(compiled_graph, payload)

    def workflow_sync(**kwargs: Any) -> Any:
        """Execute the workflow graph synchronously."""
        payload = {"inputs": kwargs, "results": {}, "messages": []}
        return asyncio.run(_run_tool_graph(compiled_graph, payload))

    return StructuredTool.from_function(
        func=workflow_sync,
        coroutine=workflow_coroutine,
        name=name,
        description=description,
        args_schema=args_schema,
    )


class WorkflowTool(BaseModel):
    """Workflow tool."""

    model_config = {"arbitrary_types_allowed": True}

    name: str
    """Name of the tool."""
    description: str
    """Description of the tool."""
    graph: SkipJsonSchema[StateGraph]
    """Workflow to be used as tool."""
    args_schema: type[BaseModel] | None = None
    """Input schema for the tool."""
    _compiled_graph: SkipJsonSchema[Runnable | None] = None
    """Cached compiled graph to avoid recompilation."""

    def get_compiled_graph(self) -> Runnable:
        """Get or compile the graph, caching the result.

        Returns:
            Compiled graph runnable
        """
        if self._compiled_graph is None:
            self._compiled_graph = self.graph.compile()
        return self._compiled_graph


@registry.register(
    NodeMetadata(
        name="AgentNode",
        description="Execute an AI agent with tools",
        category="ai",
    )
)
class AgentNode(AINode):
    """Node for executing an AI agent with tools."""

    model_config = ConfigDict(arbitrary_types_allowed=True)
    _history_key_pattern: ClassVar[re.Pattern[str]] = re.compile(r"^[A-Za-z0-9:_-]+$")
    _history_key_max_length: ClassVar[int] = 256
    _history_write_retry_limit: ClassVar[int] = 3
    _history_retry_base_backoff_seconds: ClassVar[float] = 0.025

    ai_model: str
    """Identifier of the AI chat model to use."""
    model_settings: dict | None = None
    """TODO: Implement model settings for the agent."""
    model_kwargs: dict[str, Any] = Field(
        default_factory=dict,
        description="Additional keyword arguments passed to init_chat_model.",
    )
    """Additional keyword arguments passed to init_chat_model."""
    system_prompt: str | TextTensor | None = None
    """System prompt for the agent."""
    predefined_tools: list[str] = Field(default_factory=list)
    """Tool names predefined by Orcheo."""
    workflow_tools: list[WorkflowTool] = Field(default_factory=list)
    """Workflows to be used as tools."""
    mcp_servers: dict[str, Any] = Field(default_factory=dict)
    """MCP servers to be used as tools (Connection from langchain_mcp_adapters)."""
    response_format: dict | type[BaseModel] | None = None
    """Response format for the agent."""
    max_messages: int = 30
    """Maximum number of messages to keep when sending to the agent."""
    reset_command: str = ""
    """Command that resets history. Messages before the latest reset are ignored."""
    use_graph_chat_history: bool = False
    """Enable graph-store-backed chat history loading and persistence."""
    history_namespace: list[str] = Field(default_factory=lambda: ["agent_chat_history"])
    """Namespace used for graph-store chat history items."""
    history_key_template: str = "{{conversation_key}}"
    """Template used to derive the final graph-store key."""
    history_key_candidates: list[str] = Field(
        default_factory=lambda: [
            "telegram:{{results.telegram_events_parser.chat_id}}",
            "wecom_cs:{{results.wecom_cs_sync.open_kf_id}}:{{results.wecom_cs_sync.external_userid}}",
            "wecom_aibot:{{results.wecom_ai_bot_events_parser.chat_type}}:{{results.wecom_ai_bot_events_parser.user}}",
            "wecom_dm:{{results.wecom_events_parser.user}}",
        ]
    )
    """Ordered key candidates used to resolve stable conversation identity."""
    history_value_field: str = "content"
    """Field name used when persisting text content into store records."""

    _DEFERRED_DECODE_FIELDS: ClassVar[frozenset[str]] = frozenset(
        {"history_key_template", "history_key_candidates"}
    )
    """Fields resolved by ``_resolve_history_key``, skipped in general decode."""

    def _compute_run_updates(self, state: State) -> dict[str, Any]:
        """Skip deferred fields that are resolved by _resolve_history_key."""
        updates: dict[str, Any] = {}
        for key, value in self.__dict__.items():
            if key in self._DEFERRED_DECODE_FIELDS:
                continue
            decoded = self._decode_value(value, state)
            if decoded is not value:
                updates[key] = decoded
        return updates

    def decode_variables(
        self,
        state: Any,
        *,
        config: Mapping[str, Any] | None = None,
    ) -> None:
        """Decode variables, skipping fields handled by _resolve_history_key."""
        del config
        for key, value in self.__dict__.items():
            if key in self._DEFERRED_DECODE_FIELDS:
                continue
            self.__dict__[key] = self._decode_value(value, state)

    @field_serializer("system_prompt", when_used="json")
    def _serialize_system_prompt(self, value: str | TextTensor | None) -> str | None:
        """Normalize TextTensor prompts into JSON-safe strings."""
        if isinstance(value, TextTensor):
            return value.text
        return value

    def model_post_init(self, __context: Any) -> None:
        """Normalize system prompts into TextTensor for optimization."""
        if isinstance(self.system_prompt, str):
            if "{{" in self.system_prompt and "}}" in self.system_prompt:
                return
            if (
                isinstance(self.ai_model, str)
                and ":" not in self.ai_model
                and "model_provider" not in self.model_kwargs
            ):
                return
            self.system_prompt = TextTensor(
                self.system_prompt,
                model=self.ai_model,
                model_kwargs=dict(self.model_kwargs),
            )

    def get_params(self) -> list[TextTensor]:
        """Return trainable parameters for optimizer discovery."""
        if (
            isinstance(self.system_prompt, TextTensor)
            and self.system_prompt.requires_grad
        ):
            return [self.system_prompt]
        return []

    async def _prepare_tools(self) -> list[BaseTool]:
        """Prepare the tools for the agent."""
        tools: list[BaseTool] = []

        # Resolve predefined tools from the tool registry
        for tool_name in self.predefined_tools:
            tool = tool_registry.get_tool(tool_name)
            if tool is None:
                logger.warning("Tool '%s' not found in registry, skipping", tool_name)
                continue

            # If it's already a BaseTool instance (e.g., from @tool
            # decorator), use it directly
            if isinstance(tool, BaseTool):
                tools.append(tool)
            # Otherwise, check if it's a callable factory
            elif callable(tool):
                try:
                    tool_instance = tool()
                    if not isinstance(tool_instance, BaseTool):
                        logger.error(
                            "Tool factory '%s' did not return a BaseTool instance, "
                            "got %s",
                            tool_name,
                            type(tool_instance).__name__,
                        )
                        continue
                    tools.append(tool_instance)
                except Exception as e:
                    logger.error(
                        "Failed to instantiate tool '%s': %s", tool_name, str(e)
                    )
                    continue
            else:
                logger.error(
                    "Tool '%s' is neither a BaseTool instance nor a callable factory, "
                    "got %s",
                    tool_name,
                    type(tool).__name__,
                )
                continue

        for wf_tool_def in self.workflow_tools:
            # Use cached compiled graph to avoid recompilation on every run
            compiled_graph = wf_tool_def.get_compiled_graph()

            # Create tool using factory function to properly bind variables
            # and avoid closure memory leak issues
            tool = _create_workflow_tool_func(
                compiled_graph=compiled_graph,
                name=wf_tool_def.name,
                description=wf_tool_def.description,
                args_schema=wf_tool_def.args_schema,
            )
            tools.append(tool)

        # Get MCP tools
        mcp_client = MultiServerMCPClient(connections=self.mcp_servers)
        mcp_tools = await mcp_client.get_tools()
        tools.extend(mcp_tools)

        return tools

    def _messages_from_inputs(self, inputs: Mapping[str, Any]) -> list[BaseMessage]:
        """Build LangChain messages from ChatKit-style inputs."""
        history = inputs.get("history")
        messages: list[BaseMessage] = []

        if isinstance(history, list):
            for turn in history:
                if not isinstance(turn, Mapping):
                    continue
                content = turn.get("content")
                role = turn.get("role")
                if not isinstance(content, str) or not content.strip():
                    continue
                if role == "assistant":
                    messages.append(AIMessage(content=content))
                elif role == "user":  # pragma: no branch
                    messages.append(HumanMessage(content=content))

        message_value = (
            inputs.get("message")
            or inputs.get("user_message")
            or inputs.get("query")
            or inputs.get("prompt")
        )
        if (
            isinstance(message_value, str) and message_value.strip()
        ):  # pragma: no branch
            messages.append(HumanMessage(content=message_value.strip()))

        return messages

    def _normalize_messages(self, messages: Any) -> list[BaseMessage]:
        """Normalize caller-provided messages into LangChain BaseMessages."""
        normalized: list[BaseMessage] = []
        if not isinstance(messages, list):
            return normalized

        for message in messages:
            if isinstance(message, BaseMessage):
                normalized.append(message)
                continue
            if not isinstance(message, Mapping):
                continue
            content = message.get("content")
            role = message.get("role")
            if not isinstance(content, str) or not content.strip():
                continue
            if role == "assistant":
                normalized.append(AIMessage(content=content))
            else:
                normalized.append(HumanMessage(content=content))

        return normalized

    def _apply_reset_command(self, messages: list[BaseMessage]) -> list[BaseMessage]:
        """Trim messages to start from the latest reset command (inclusive)."""
        if not self.reset_command:
            return messages
        for i in range(len(messages) - 1, -1, -1):
            content = messages[i].content
            if isinstance(content, str) and content.strip() == self.reset_command:
                return messages[i:]
        return messages

    def _build_messages(
        self,
        state: State,
        config: RunnableConfig | None = None,
    ) -> list[BaseMessage]:
        """Construct the message list for the agent invocation."""
        existing_messages = self._normalize_messages(state.get("messages"))
        if existing_messages:
            messages = self._apply_reset_command(existing_messages)

            configurable = (
                config.get("configurable", {}) if isinstance(config, Mapping) else {}
            )
            has_checkpointer = (
                isinstance(configurable, Mapping)
                and "thread_id" in configurable
                and "__pregel_checkpointer" in configurable
            )
            if has_checkpointer:
                inputs = state.get("inputs", {}) if isinstance(state, Mapping) else {}
                new_messages = self._messages_from_inputs(inputs)
                if new_messages:
                    messages.extend(new_messages)

            return messages[-self.max_messages :]

        inputs = state.get("inputs", {}) if isinstance(state, Mapping) else {}
        messages = self._messages_from_inputs(inputs)
        messages = self._apply_reset_command(messages)
        return messages[-self.max_messages :]

    def _get_graph_store(self, config: RunnableConfig | None) -> Any | None:
        """Return the runtime graph store when available."""
        return _get_graph_store_fn(config)

    def _history_namespace_tuple(self) -> tuple[str, ...]:
        namespace = tuple(
            entry.strip()
            for entry in self.history_namespace
            if isinstance(entry, str) and entry.strip()
        )
        return namespace or ("agent_chat_history",)

    def _validate_history_key(self, key: str) -> tuple[str | None, str]:
        """Validate a candidate history key and return failure reason when invalid."""
        candidate = key.strip()
        if not candidate:
            return None, "empty"
        if "{{" in candidate or "}}" in candidate:
            return None, "unresolved_template"
        if len(candidate) > self._history_key_max_length:
            return None, "too_long"
        if self._history_key_pattern.fullmatch(candidate) is None:
            return None, "invalid_chars"
        return candidate, "ok"

    def _resolve_history_key(
        self,
        state: State,
        config: RunnableConfig | None,
    ) -> str | None:
        """Resolve and validate the final graph-store history key."""
        del config
        conversation_key: str | None = None

        for candidate in self.history_key_candidates:
            if not isinstance(candidate, str):
                continue
            resolved = self._decode_string_value(candidate, state)
            rendered = (
                str(resolved).strip()
                if isinstance(resolved, str | int | float | bool)
                else candidate
            )
            valid_key, status = self._validate_history_key(rendered)
            if valid_key is not None:
                conversation_key = valid_key
                break
            if rendered:
                logger.debug(
                    "AgentNode '%s' rejected history key candidate '%s' (%s).",
                    self.name,
                    rendered,
                    status,
                )

        if conversation_key is None:
            logger.warning(
                "AgentNode '%s' skipped graph history: no valid conversation key "
                "resolved.",
                self.name,
            )
            return None

        history_template_state = cast(
            State,
            {
                **(dict(state) if isinstance(state, Mapping) else {}),
                "conversation_key": conversation_key,
            },
        )
        rendered_key_value = self._decode_string_value(
            self.history_key_template,
            history_template_state,
        )
        rendered_key = (
            str(rendered_key_value).strip()
            if isinstance(rendered_key_value, str)
            else self.history_key_template.strip()
        )
        final_key, status = self._validate_history_key(rendered_key)
        if final_key is None:
            logger.warning(
                "AgentNode '%s' skipped graph history: resolved key '%s' is invalid "
                "(%s).",
                self.name,
                rendered_key,
                status,
            )
            return None

        return final_key

    async def _store_get_item(
        self,
        store: Any,
        namespace: tuple[str, ...],
        key: str,
    ) -> Any | None:
        """Read a single store item using async/sync API variants."""
        aget = getattr(store, "aget", None)
        if callable(aget):
            return await aget(namespace, key)

        get = getattr(store, "get", None)
        if callable(get):
            result = get(namespace, key)
            if asyncio.iscoroutine(result):
                return await result
            return result
        return None

    async def _store_put_item(
        self,
        store: Any,
        namespace: tuple[str, ...],
        key: str,
        value: dict[str, Any],
    ) -> None:
        """Write a single store item using async/sync API variants."""
        aput = getattr(store, "aput", None)
        if callable(aput):
            await aput(namespace, key, value)
            return

        put = getattr(store, "put", None)
        if callable(put):
            result = put(namespace, key, value)
            if asyncio.iscoroutine(result):
                await result

    def _history_payload_from_item(self, item: Any | None) -> Mapping[str, Any]:
        """Extract history payload from a LangGraph item-like response."""
        if item is None:
            return {}
        if isinstance(item, Mapping):
            payload = item.get("value")
            return payload if isinstance(payload, Mapping) else {}
        value = getattr(item, "value", None)
        return value if isinstance(value, Mapping) else {}

    def _normalize_history_store_messages(self, payload: Any) -> list[BaseMessage]:
        """Normalize persisted history payload into LangChain messages."""
        normalized: list[BaseMessage] = []
        if not isinstance(payload, list):
            return normalized

        for item in payload:
            if not isinstance(item, Mapping):
                continue
            role = item.get("role")
            content = item.get(self.history_value_field, item.get("content"))
            if not isinstance(content, str) or not content.strip():
                continue
            if role == "assistant":
                normalized.append(AIMessage(content=content))
            elif role == "user":  # pragma: no branch
                normalized.append(HumanMessage(content=content))
        return normalized

    def _serialize_history_messages(
        self, messages: list[BaseMessage]
    ) -> list[dict[str, str]]:
        """Serialize user/assistant messages for graph-store persistence."""
        serialized: list[dict[str, str]] = []
        for message in messages:
            content = message.content
            text = content if isinstance(content, str) else str(content)
            if not text.strip():
                continue
            if isinstance(message, AIMessage):
                role = "assistant"
            elif isinstance(message, HumanMessage):
                role = "user"
            else:
                continue
            payload = {"role": role, self.history_value_field: text}
            if self.history_value_field != "content":
                payload["content"] = text
            serialized.append(payload)
        return serialized

    def _filter_user_assistant_messages(
        self,
        messages: list[BaseMessage],
    ) -> list[BaseMessage]:
        return [
            message
            for message in messages
            if isinstance(message, HumanMessage | AIMessage)
        ]

    def _message_signature(self, message: BaseMessage) -> tuple[str, str]:
        """Create a deterministic identity for message overlap detection."""
        content = message.content
        text = content if isinstance(content, str) else str(content)
        if isinstance(message, AIMessage):
            return "assistant", text
        if isinstance(message, HumanMessage):
            return "user", text
        return message.type, text

    def _starts_with_messages(
        self,
        messages: list[BaseMessage],
        prefix: list[BaseMessage],
    ) -> bool:
        if len(prefix) > len(messages):
            return False
        return all(
            self._message_signature(messages[index])
            == self._message_signature(prefix[index])
            for index in range(len(prefix))
        )

    def _suffix_prefix_overlap(
        self,
        existing: list[BaseMessage],
        observed: list[BaseMessage],
    ) -> int:
        """Return overlap size between existing suffix and observed prefix."""
        max_overlap = min(len(existing), len(observed))
        for overlap in range(max_overlap, 0, -1):
            existing_slice = existing[-overlap:]
            observed_slice = observed[:overlap]
            if all(
                self._message_signature(existing_slice[index])
                == self._message_signature(observed_slice[index])
                for index in range(overlap)
            ):
                return overlap
        return 0

    def _history_version_from_payload(self, payload: Mapping[str, Any]) -> int:
        raw_version = payload.get("version", 0)
        if isinstance(raw_version, int) and raw_version >= 0:
            return raw_version
        if isinstance(raw_version, str) and raw_version.isdigit():
            return int(raw_version)
        return 0

    def _extract_observed_messages(
        self,
        current_messages: list[BaseMessage],
        inference_messages: list[BaseMessage],
        result: Mapping[str, Any],
    ) -> list[BaseMessage]:
        """Collect user/assistant turns observed during this run."""
        current_turns = self._filter_user_assistant_messages(current_messages)
        result_messages = self._filter_user_assistant_messages(
            self._normalize_messages(result.get("messages"))
        )

        delta = result_messages
        if self._starts_with_messages(result_messages, inference_messages):
            delta = result_messages[len(inference_messages) :]
        elif self._starts_with_messages(result_messages, current_turns):
            delta = result_messages[len(current_turns) :]

        return current_turns + delta

    async def _load_graph_history_messages(
        self,
        *,
        store: Any,
        namespace: tuple[str, ...],
        key: str,
    ) -> list[BaseMessage]:
        """Read full persisted history for the given namespace/key."""
        item = await self._store_get_item(store, namespace, key)
        payload = self._history_payload_from_item(item)
        return self._normalize_history_store_messages(payload.get("messages"))

    async def _persist_graph_history(
        self,
        *,
        store: Any,
        namespace: tuple[str, ...],
        key: str,
        observed_messages: list[BaseMessage],
    ) -> None:
        """Append observed turns to store history using bounded retry."""
        observed = self._filter_user_assistant_messages(observed_messages)
        if not observed:
            return

        for attempt in range(self._history_write_retry_limit):
            try:
                existing_item = await self._store_get_item(store, namespace, key)
            except Exception:
                logger.warning(
                    "AgentNode '%s' failed to read graph history before write "
                    "(key='%s').",
                    self.name,
                    key,
                )
                return

            existing_payload = self._history_payload_from_item(existing_item)
            existing_messages = self._normalize_history_store_messages(
                existing_payload.get("messages")
            )
            overlap = self._suffix_prefix_overlap(existing_messages, observed)
            new_messages = observed[overlap:]
            if not new_messages:
                return

            merged_messages = existing_messages + new_messages
            next_version = self._history_version_from_payload(existing_payload) + 1
            payload = {
                "version": next_version,
                "messages": self._serialize_history_messages(merged_messages),
            }

            try:
                await self._store_put_item(store, namespace, key, payload)
                written_item = await self._store_get_item(store, namespace, key)
            except Exception:
                if attempt + 1 >= self._history_write_retry_limit:
                    logger.warning(
                        "AgentNode '%s' failed to persist graph history after %d "
                        "attempts (key='%s').",
                        self.name,
                        self._history_write_retry_limit,
                        key,
                    )
                    return
                backoff = (2**attempt) * self._history_retry_base_backoff_seconds
                await asyncio.sleep(backoff + random.uniform(0.0, 0.01))
                continue

            written_payload = self._history_payload_from_item(written_item)
            written_messages = self._normalize_history_store_messages(
                written_payload.get("messages")
            )
            written_version = self._history_version_from_payload(written_payload)
            if (
                written_version >= next_version
                and len(written_messages) >= len(merged_messages)
                and self._starts_with_messages(
                    written_messages[-len(merged_messages) :],
                    merged_messages,
                )
            ):
                return

            if attempt + 1 >= self._history_write_retry_limit:
                logger.warning(
                    "AgentNode '%s' detected persistent graph history write conflicts "
                    "after %d attempts (key='%s').",
                    self.name,
                    self._history_write_retry_limit,
                    key,
                )
                return

            backoff = (2**attempt) * self._history_retry_base_backoff_seconds
            await asyncio.sleep(backoff + random.uniform(0.0, 0.01))

    async def run(self, state: State, config: RunnableConfig) -> dict[str, Any]:
        """Execute the agent and return results."""
        tools = await self._prepare_tools()

        response_format_strategy = None
        if self.response_format is not None:
            response_format_strategy = ProviderStrategy(self.response_format)  # type: ignore[arg-type]

        # Initialize chat model with model_kwargs
        model = init_chat_model(self.ai_model, **self.model_kwargs)
        agent = create_agent(
            model,
            tools=tools,
            system_prompt=self._system_prompt_text,
            response_format=response_format_strategy,
        )
        # TODO: for models that don't support ProviderStrategy, use ToolStrategy

        current_messages = self._build_messages(state, config)
        messages = list(current_messages)

        history_store: Any | None = None
        history_namespace: tuple[str, ...] = ()
        history_key: str | None = None

        if self.use_graph_chat_history:
            history_store = self._get_graph_store(config)
            history_key = self._resolve_history_key(state, config)
            history_namespace = self._history_namespace_tuple()
            if history_store is None:
                logger.warning(
                    "AgentNode '%s' enabled graph history but runtime store is "
                    "missing.",
                    self.name,
                )
            elif history_key is not None:
                try:
                    persisted_messages = await self._load_graph_history_messages(
                        store=history_store,
                        namespace=history_namespace,
                        key=history_key,
                    )
                except Exception:
                    logger.warning(
                        "AgentNode '%s' failed to read graph history (key='%s'); "
                        "falling back to in-memory messages.",
                        self.name,
                        history_key,
                    )
                else:
                    messages = persisted_messages + current_messages
                    if len(messages) > self.max_messages:
                        logger.info(
                            "AgentNode '%s' truncated merged history from %d to %d "
                            "messages (key='%s').",
                            self.name,
                            len(messages),
                            self.max_messages,
                            history_key,
                        )
                    messages = messages[-self.max_messages :]

        # Execute agent with normalized messages as input
        payload: dict[str, Any] = {"messages": messages}
        with tool_execution_context(config):
            result = await agent.ainvoke(payload, config)  # type: ignore[arg-type]

        if (
            self.use_graph_chat_history
            and history_store is not None
            and history_key is not None
            and isinstance(result, Mapping)
        ):
            observed_messages = self._extract_observed_messages(
                current_messages,
                messages,
                result,
            )
            await self._persist_graph_history(
                store=history_store,
                namespace=history_namespace,
                key=history_key,
                observed_messages=observed_messages,
            )
        return result

    @property
    def _system_prompt_text(self) -> str | None:
        if isinstance(self.system_prompt, TextTensor):
            return self.system_prompt.text
        return self.system_prompt


@registry.register(
    NodeMetadata(
        name="AgentReplyExtractorNode",
        description="Extract the final assistant reply from agent messages",
        category="ai",
    )
)
class AgentReplyExtractorNode(TaskNode):
    """Extract the last assistant message from the agent output.

    After an :class:`AgentNode` runs, the workflow state contains a
    ``messages`` list mixing user and assistant turns.  This node scans
    that list in reverse and returns the most recent assistant reply as
    plain text.
    """

    fallback_message: str = Field(
        default="Sorry, something went wrong. Please try again later.",
        description="Message returned when no assistant reply is found",
    )

    async def run(self, state: State, config: RunnableConfig) -> dict[str, Any]:
        """Return ``{"agent_reply": "..."}`` from the last AI message."""
        messages = state.get("messages", [])
        for msg in reversed(messages):
            if isinstance(msg, dict):
                if msg.get("role") == "assistant":
                    content = msg.get("content", "")
                    if content:
                        return {"agent_reply": str(content)}
            elif isinstance(msg, BaseMessage) and msg.type == "ai" and msg.content:
                content = msg.content
                return {
                    "agent_reply": content if isinstance(content, str) else str(content)
                }
        return {"agent_reply": self.fallback_message}


@registry.register(
    NodeMetadata(
        name="LLMNode",
        description="Execute a text-only LLM call",
        category="ai",
    )
)
class LLMNode(AgentNode):
    """Node for executing an LLM on a single text input."""

    input_text: str | None = None
    """Text input to be processed by the LLM."""
    instruction: str | None = None
    """Optional instruction for post-processing the input."""
    user_message: str | None = None
    """Optional user message for language or tone inference."""
    draft_reply: str | None = None
    """Draft reply to be post-processed."""

    async def run(self, state: State, config: RunnableConfig) -> dict[str, Any]:
        """Execute the LLM with a single text input."""
        messages = self._build_messages(state)
        if not messages:
            return {"messages": []}

        tools = await self._prepare_tools()

        response_format_strategy = None
        if self.response_format is not None:
            response_format_strategy = ProviderStrategy(self.response_format)  # type: ignore[arg-type]

        model = init_chat_model(self.ai_model, **self.model_kwargs)
        agent = create_agent(
            model,
            tools=tools,
            system_prompt=self._system_prompt_text,
            response_format=response_format_strategy,
        )

        payload: dict[str, Any] = {"messages": messages}
        with tool_execution_context(config):
            result = await agent.ainvoke(payload, config)  # type: ignore[arg-type]
        return result

    def _build_messages(
        self,
        _state: State,
        config: RunnableConfig | None = None,
    ) -> list[BaseMessage]:
        """Construct a single-turn message list for the LLM."""
        draft_reply = self._normalize_text(self.draft_reply)
        input_text = self._normalize_text(self.input_text)
        if not draft_reply and not input_text:
            return []

        base_text = draft_reply or input_text
        user_message = self._normalize_text(self.user_message)

        if user_message:
            content_text = f"User message:\n{user_message}\n\nDraft reply:\n{base_text}"
        else:
            content_text = base_text

        instruction = self._normalize_text(self.instruction)
        if instruction:
            content = f"Instruction:\n{instruction}\n\nText:\n{content_text}"
        else:
            content = content_text
        return [HumanMessage(content=content)]

    @staticmethod
    def _normalize_text(value: str | None) -> str:
        if isinstance(value, str):
            return value.strip()
        return ""
