"""Deep research agent node for autonomous multi-step research workflows."""

from __future__ import annotations
import logging
from typing import Any
from deepagents import create_deep_agent
from langchain.agents.structured_output import ProviderStrategy
from langchain.chat_models import init_chat_model
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import BaseTool
from langchain_mcp_adapters.client import MultiServerMCPClient
from pydantic import ConfigDict, Field
from orcheo.graph.state import State
from orcheo.nodes.agent_tools.context import tool_execution_context
from orcheo.nodes.agent_tools.registry import tool_registry
from orcheo.nodes.ai import WorkflowTool, _create_workflow_tool_func
from orcheo.nodes.base import AINode
from orcheo.nodes.registry import NodeMetadata, registry
from orcheo.skills.manager import SkillManager
from orcheo.skills.paths import get_skills_dir


logger = logging.getLogger(__name__)


@registry.register(
    NodeMetadata(
        name="DeepAgentNode",
        description=(
            "Execute an autonomous deep-research agent with configurable "
            "iteration depth for multi-step tool use and synthesis"
        ),
        category="ai",
    )
)
class DeepAgentNode(AINode):
    """Node for executing a deep-research agent with multi-step iteration.

    Wraps ``create_deep_agent`` from LangChain's ``deepagents`` package,
    which provides built-in planning tools, a virtual file-system backend,
    sub-agent spawning, and context summarisation middleware.

    Unlike :class:`AgentNode` which includes chat history management and
    multi-turn conversation tracking, ``DeepAgentNode`` is focused on
    single-invocation deep research: the agent receives a query, plans
    its approach, iterates over tools up to ``max_iterations`` times, and
    returns a synthesised result.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    ai_model: str
    """Identifier of the AI chat model to use (e.g. ``"openai:gpt-4o"``)."""
    system_prompt: str | None = None
    """Custom system instructions prepended to the deep-agent base prompt."""
    research_prompt: str | None = None
    """Research-specific instructions appended to ``system_prompt``."""
    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."""
    max_iterations: int = Field(
        default=100,
        description="Maximum number of agent iterations (recursion limit).",
    )
    """Maximum number of agent iterations (recursion limit)."""
    model_kwargs: dict[str, Any] = Field(default_factory=dict)
    """Additional keyword arguments passed to ``init_chat_model``."""
    response_format: dict | None = None
    """Structured output response format for the agent."""
    input_query: str | None = None
    """Direct text query for the agent. Supports variable interpolation."""
    skills: list[str] | None = None
    """Skill source paths loaded by the deep-agent middleware."""
    memory: list[str] | None = None
    """Memory file paths (``AGENTS.md``) loaded at agent startup."""
    debug: bool = False
    """Enable deep-agent debug mode."""

    def _build_system_prompt(self) -> str | None:
        """Build the combined system prompt from base and research prompts.

        Returns:
            Combined prompt string or None if both are empty.
        """
        parts: list[str] = []
        if self.system_prompt:
            parts.append(self.system_prompt.strip())
        if self.research_prompt:
            parts.append(self.research_prompt.strip())
        return "\n\n".join(parts) if parts else None

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

        Resolves predefined tools from the registry, compiles workflow tools,
        and loads MCP server tools.

        Returns:
            List of all resolved tools.
        """
        tools: list[BaseTool] = []

        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 isinstance(tool, BaseTool):
                tools.append(tool)
            elif callable(tool):
                try:
                    tool_instance = tool()
                    if not isinstance(tool_instance, BaseTool):
                        logger.error(
                            "Tool factory '%s' did not return a BaseTool, got %s",
                            tool_name,
                            type(tool_instance).__name__,
                        )
                        continue
                    tools.append(tool_instance)
                except Exception as exc:
                    logger.error(
                        "Failed to instantiate tool '%s': %s", tool_name, str(exc)
                    )
                    continue
            else:
                logger.error(
                    "Tool '%s' is not a BaseTool or callable, got %s",
                    tool_name,
                    type(tool).__name__,
                )
                continue

        for wf_tool_def in self.workflow_tools:
            compiled_graph = wf_tool_def.get_compiled_graph()
            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,
                output_path=wf_tool_def.output_path,
            )
            tools.append(tool)

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

        return tools

    def _build_messages(self, state: State) -> list[BaseMessage]:
        """Build the message list for the agent invocation.

        Uses ``input_query`` if set, otherwise falls back to extracting
        the message from workflow inputs.

        Returns:
            List of messages to send to the agent.
        """
        if self.input_query:
            return [HumanMessage(content=self.input_query)]

        inputs = state.get("inputs", {})
        if isinstance(inputs, dict):
            for key in ("query", "message", "prompt", "input"):
                value = inputs.get(key)
                if isinstance(value, str) and value.strip():
                    return [HumanMessage(content=value.strip())]

        return []

    def _resolve_skills(self) -> list[str] | None:
        """Resolve skill paths for the agent.

        When ``self.skills`` is ``None``, automatically discovers all
        installed skills from the Orcheo skills directory.  When set
        explicitly, returns the configured list unchanged.

        Returns:
            List of skill directory paths, or ``None`` if no skills
            are available.
        """
        if self.skills is not None:
            return self.skills

        try:
            manager = SkillManager(skills_dir=get_skills_dir())
            paths = manager.get_installed_skill_paths()
        except Exception:
            logger.warning("Failed to discover installed skills", exc_info=True)
            return None

        return paths if paths else None

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

        Creates a deep agent via ``create_deep_agent`` from the
        ``deepagents`` package, which includes built-in planning,
        file-system, sub-agent, and summarisation middleware.

        Args:
            state: Current workflow state.
            config: Runnable configuration.

        Returns:
            Agent execution result with messages.
        """
        tools = await self._prepare_tools()

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

        model_kwargs = self.model_kwargs
        if model_kwargs:
            model = init_chat_model(self.ai_model, **model_kwargs)
        else:
            model = self.ai_model  # type: ignore[assignment]

        combined_prompt = self._build_system_prompt()
        resolved_skills = self._resolve_skills()
        agent = create_deep_agent(
            model,
            tools=tools,
            system_prompt=combined_prompt,
            response_format=response_format_strategy,
            skills=resolved_skills,
            memory=self.memory,
            debug=self.debug,
        )

        messages = self._build_messages(state)
        payload: dict[str, Any] = {"messages": messages}

        with tool_execution_context(config):
            result = await agent.ainvoke(
                payload,  # type: ignore[arg-type]
                config={
                    **(config or {}),
                    "recursion_limit": self.max_iterations,
                },
            )

        return result


__all__ = ["DeepAgentNode"]
