"""ReasoningAgentMacro - Multi-step reasoning with adaptive tool calling.

Architecture:
- **Native Tool Calling**: Uses LLM adapter's aresponse_with_tools for OpenAI/Anthropic/Gemini
- **Text-Based Fallback**: INVOKE_TOOL: parsing for adapters without native support
- **Runtime Detection**: Adapts automatically based on adapter capabilities
- **Fallback Policy Support**: Seamlessly handles adapter switching during failures

This adaptive approach ensures:
- Optimal performance with native tool calling when available
- Compatibility with any LLM via text-based fallback
- Seamless integration with hexDAG's fallback policies
- Single graph works with multiple adapter types

Example workflow:
1. LLM reasoning step → tries native tools first, falls back to text if needed
2. Tool executor → executes parsed tool calls
3. Result merger → combines reasoning with tool results
4. Next step → continues with context from previous tools
"""

from typing import Any

from pydantic import field_validator

from hexdag.builtin.nodes.function_node import FunctionNode
from hexdag.builtin.nodes.llm_node import LLMNode
from hexdag.builtin.nodes.tool_utils import (
    ToolCallFormat,
    ToolDefinition,
    ToolParser,
    ToolSchemaConverter,
)
from hexdag.builtin.prompts.tool_prompts import get_tool_prompt_for_format
from hexdag.core.configurable import ConfigurableMacro, MacroConfig
from hexdag.core.context import get_port
from hexdag.core.domain.dag import DirectedGraph
from hexdag.core.logging import get_logger
from hexdag.core.orchestration.prompt import PromptTemplate
from hexdag.core.ports.llm import Message, MessageList
from hexdag.core.resolver import resolve_function

logger = get_logger(__name__)


class ReasoningAgentConfig(MacroConfig):
    """Configuration for ReasoningAgentMacro.

    Attributes
    ----------
    main_prompt : str
        Primary prompt for reasoning
    max_steps : int
        Maximum reasoning iterations (default: 3)
    allowed_tools : list[str]
        Tool names available to the agent (e.g., ["core:search", "core:calculate"])
        Uses qualified names (namespace:name)
    tool_format : ToolCallFormat
        Tool calling format for text-based fallback: FUNCTION_CALL, JSON, or MIXED (default: MIXED)
        Only used when adapter doesn't support native tool calling
    """

    main_prompt: str
    max_steps: int = 5
    allowed_tools: list[str] = []
    tool_format: ToolCallFormat = ToolCallFormat.MIXED

    @field_validator("main_prompt", mode="before")
    @classmethod
    def convert_prompt_input(cls, v: Any) -> str:
        """Convert PromptInput to string."""
        if isinstance(v, str):
            return v
        if isinstance(v, PromptTemplate):
            return v.template
        if hasattr(v, "template"):
            return str(v.template)
        return str(v)


class ReasoningAgentMacro(ConfigurableMacro):
    """Multi-step reasoning agent with adaptive tool calling.

    Automatically detects and adapts to LLM adapter capabilities:
    - **Native mode**: OpenAI/Anthropic/Gemini → clean prompts, structured tool calls
    - **Text mode**: Local/other LLMs → INVOKE_TOOL: parsing from text
    - **Runtime adaptive**: Checks hasattr(llm, 'aresponse_with_tools') at execution time

    This adaptive approach enables:
    1. **Fallback Policy Support**: When a node fails and policy switches adapters,
       the same graph works with the new adapter type
    2. **Flexibility**: Single graph handles multiple adapter types
    3. **Optimal Performance**: Uses native tools when available

    Architecture per reasoning step:
    ```
    LLM Node (adaptive) → Tool Executor → Result Merger
    ```

    The LLM node adapts at runtime:
    - Tries native tool calling first (if available)
    - Falls back to text-based parsing automatically
    - Both paths prepared, chosen based on adapter capabilities
    """

    Config = ReasoningAgentConfig

    def expand(
        self,
        instance_name: str,
        inputs: dict[str, Any],
        dependencies: list[str],
    ) -> DirectedGraph:
        """Expand into reasoning chain with adaptive tool calling strategy."""
        graph = DirectedGraph()
        config: ReasoningAgentConfig = self.config  # type: ignore[assignment]

        fn_factory = FunctionNode()

        # Build tool schemas for native calling
        tool_schemas = self._build_tool_schemas_for_native(config.allowed_tools)

        # Build tool list for text-based fallback
        tool_list_text = self._build_tool_list_for_text(config.allowed_tools)
        tool_prompt = get_tool_prompt_for_format(config.tool_format)

        # Build reasoning chain
        prev_step: str | None = None
        for step_idx in range(config.max_steps):
            step_name_prefix = f"{instance_name}_step_{step_idx}"
            step_deps: list[str] = dependencies if step_idx == 0 else [prev_step]  # type: ignore[list-item]

            # Create LLM subgraph (unified LLMNode + adapter)
            llm_subgraph = self._create_llm_subgraph(
                fn_factory,
                step_name_prefix,
                step_idx,
                prev_step,
                config,
                tool_schemas,
                tool_list_text,
                tool_prompt,
                step_deps,
            )
            graph |= llm_subgraph

            # Create tool executor node
            tool_executor = self._create_tool_executor_node(fn_factory, step_name_prefix, config)
            graph += tool_executor

            # Create result merger node
            result_merger = self._create_result_merger_node(fn_factory, step_name_prefix)
            graph += result_merger

            prev_step = f"{step_name_prefix}_result_merger"

        # Final consolidation using unified LLMNode
        if prev_step is None:
            raise ValueError("prev_step is None")

        llm_factory = LLMNode()
        final_llm = llm_factory(
            name=f"{instance_name}_final",
            prompt_template=f"""All reasoning steps and tool results:
{{{{{prev_step}}}}}

Provide your final conclusion based on all reasoning and evidence gathered.""",
            deps=[prev_step],
        )
        graph += final_llm

        return graph

    def _create_llm_subgraph(
        self,
        fn_factory: FunctionNode,
        step_name: str,
        step_idx: int,
        prev_step: str | None,
        config: ReasoningAgentConfig,
        tool_schemas: list[dict[str, Any]],
        tool_list_text: str,
        tool_prompt: Any,
        deps: list[str],
    ) -> DirectedGraph:
        """Create LLM nodes using unified LLMNode architecture.

        Returns a subgraph with:
        - LLMNode: unified node for prompt + LLM call
        - FunctionNode: adapts response and parses tool calls
        """
        subgraph = DirectedGraph()

        # Build prompt content
        if step_idx == 0:
            base_prompt = config.main_prompt
        else:
            base_prompt = f"""Previous reasoning and tool results:
{{{{{prev_step}}}}}

Continue reasoning. Use tools if needed to gather more information."""

        # Add tool instructions for text-based fallback
        full_prompt = f"""{base_prompt}

## Available Tools
{tool_list_text}

{tool_prompt.template if hasattr(tool_prompt, "template") else str(tool_prompt)}"""

        # Node 1: Unified LLMNode for prompt + LLM call
        llm_factory = LLMNode()
        llm_node = llm_factory(
            name=f"{step_name}_llm",
            prompt_template=full_prompt,
            deps=deps,
        )
        subgraph += llm_node

        # Node 2: Adapter function to normalize response and parse tool calls
        async def normalize_response(input_data: str, **kwargs: Any) -> dict[str, Any]:
            """Normalize LLM response to unified format with tool parsing."""
            text = input_data if isinstance(input_data, str) else str(input_data)

            # Try native tool calling if LLM port supports it
            llm_port = get_port("llm")
            has_tool_support = (
                llm_port and hasattr(llm_port, "aresponse_with_tools") and tool_schemas
            )

            if has_tool_support:
                # Check if adapter actually implements aresponse_with_tools
                for cls in llm_port.__class__.__mro__:
                    if (
                        cls.__name__ not in ["LLM", "Protocol", "object"]
                        and "aresponse_with_tools" in cls.__dict__
                    ):
                        # Native tool calling available but we already got text response
                        # Parse any tool calls from the text (in case it includes them)
                        break

            # Text-based - parse tool calls from text
            parsed_calls = ToolParser.parse_tool_calls(text, format=config.tool_format)

            return {
                "content": text,
                "tool_calls": [
                    {"id": f"text_{i}", "name": tc.name, "arguments": tc.params}
                    for i, tc in enumerate(parsed_calls)
                ],
                "strategy": "text",
            }

        adapter_node = fn_factory(
            name=f"{step_name}_adapter",
            fn=normalize_response,
            deps=[f"{step_name}_llm"],
        )
        subgraph += adapter_node

        return subgraph

    def _create_tool_executor_node(
        self, fn_factory: FunctionNode, step_name: str, config: ReasoningAgentConfig
    ) -> Any:
        """Create node that executes tool calls (from native or parsed)."""

        async def execute_tools(
            input_data: Any, tool_router: Any = None, **kwargs: Any
        ) -> dict[str, Any]:
            """Execute tool calls and return results."""
            llm_output = input_data
            tool_calls = llm_output.get("tool_calls", [])

            if not tool_calls:
                return {
                    "llm_content": llm_output.get("content", ""),
                    "tool_results": [],
                    "has_tools": False,
                }

            # Build tool name mapping for allowed tools
            tool_name_map = {tool: tool for tool in config.allowed_tools}

            # Execute tools
            tool_results = []
            for tc in tool_calls:
                try:
                    tool_name = tc["name"]
                    resolved_name = tool_name_map.get(tool_name) or tool_name

                    # Execute tool
                    if tool_router:
                        result = await tool_router.acall_tool(resolved_name, tc["arguments"])
                    else:
                        # Direct resolution via module path
                        import asyncio

                        tool = resolve_function(resolved_name)
                        if asyncio.iscoroutinefunction(tool):
                            result = await tool(**tc["arguments"])
                        else:
                            result = tool(**tc["arguments"])

                    tool_results.append({
                        "tool_name": tc["name"],
                        "arguments": tc["arguments"],
                        "result": result,
                    })
                except Exception as e:
                    logger.error(f"Tool execution error for {tc['name']}: {e}")
                    tool_results.append({
                        "tool_name": tc["name"],
                        "arguments": tc["arguments"],
                        "error": str(e),
                    })

            return {
                "llm_content": llm_output.get("content", ""),
                "tool_results": tool_results,
                "has_tools": True,
            }

        return fn_factory(
            name=f"{step_name}_tool_executor",
            fn=execute_tools,
            deps=[f"{step_name}_adapter"],
        )

    def _create_result_merger_node(self, fn_factory: FunctionNode, step_name: str) -> Any:
        """Create node that merges LLM reasoning with tool results."""

        async def merge_results(input_data: Any, **kwargs: Any) -> str:
            """Combine LLM content with tool results into readable format."""
            executor_output = input_data
            llm_content: str = executor_output.get("llm_content", "")
            tool_results = executor_output.get("tool_results", [])
            has_tools = executor_output.get("has_tools", False)

            if not has_tools:
                return llm_content

            # Format tool results
            results_text = "\n\n## Tool Execution Results:\n"
            for tr in tool_results:
                if "error" in tr:
                    results_text += f"- {tr['tool_name']}: ERROR - {tr['error']}\n"
                else:
                    results_text += f"- {tr['tool_name']}: {tr['result']}\n"

            return f"{llm_content}{results_text}"

        return fn_factory(
            name=f"{step_name}_result_merger", fn=merge_results, deps=[f"{step_name}_tool_executor"]
        )

    def _create_final_consolidation_fn(self, prev_step: str) -> Any:
        """Create function for final consolidation of all reasoning."""

        async def consolidate(input_data: Any, llm: Any, **kwargs: Any) -> str:
            """Consolidate all reasoning steps into final answer."""
            all_reasoning = kwargs.get(prev_step, "")

            messages = MessageList([
                Message(
                    role="user",
                    content=f"""All reasoning steps and tool results:
{all_reasoning}

Provide your final conclusion based on all reasoning and evidence gathered.""",
                )
            ])

            response = await llm.aresponse(messages)
            return response or ""

        return consolidate

    def _build_tool_schemas_for_native(self, allowed_tools: list[str]) -> list[dict[str, Any]]:
        """Build OpenAI-format tool schemas for native calling."""
        schemas = []
        for tool_name in allowed_tools:
            try:
                # Resolve tool function to get its docstring
                tool_fn = resolve_function(tool_name)
                description = tool_fn.__doc__ or f"Tool {tool_name}"
                # Take first line of docstring
                description = description.split("\n")[0].strip()

                # Build ToolDefinition
                tool_def = ToolDefinition(
                    name=tool_name,
                    simplified_description=description,
                    detailed_description=description,
                    parameters=[],
                    examples=[],
                )

                # Convert to OpenAI format
                schema = ToolSchemaConverter.to_openai_schema(tool_def)
                schemas.append(schema)
            except Exception as e:
                logger.warning(f"Could not build schema for tool {tool_name}: {e}")

        return schemas

    def _build_tool_list_for_text(self, allowed_tools: list[str]) -> str:
        """Build text-format tool list for fallback mode."""
        if not allowed_tools:
            return "No tools available"

        tool_lines = []
        for tool_name in allowed_tools:
            try:
                tool_fn = resolve_function(tool_name)
                description = tool_fn.__doc__ or "No description"
                # Take first line of docstring
                description = description.split("\n")[0].strip()
                tool_lines.append(f"  - {tool_name}: {description}")
            except Exception:
                tool_lines.append(f"  - {tool_name}: Tool description unavailable")

        return "\n".join(tool_lines)
