"""Graph node base class and built-in node types."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from thryve.core.agent.agent import Agent
    from thryve.core.backends.adapter import BackendAdapter
    from thryve.core.context.models import Message
    from thryve.core.tools.models import Tool


class Node(ABC):
    """Base class for a DAG workflow node."""

    def __init__(self, name: str, description: str = "") -> None:
        self.name = name
        self.description = description

    @abstractmethod
    async def execute(self, inputs: dict[str, Any]) -> Any:
        """Execute the node's logic.

        Parameters
        ----------
        inputs:
            A dict mapping upstream node names to their outputs, plus any
            initial inputs provided at graph execution time.
        """

    def __repr__(self) -> str:
        return f"{type(self).__name__}(name={self.name!r})"


# ---------------------------------------------------------------------------
# Concrete node types
# ---------------------------------------------------------------------------

class FunctionNode(Node):
    """Wraps an arbitrary (async or sync) callable as a graph node."""

    def __init__(self, name: str, func: Any, description: str = "") -> None:
        super().__init__(name, description)
        self._func = func

    async def execute(self, inputs: dict[str, Any]) -> Any:
        import asyncio
        import inspect

        if inspect.iscoroutinefunction(self._func):
            return await self._func(inputs)
        return await asyncio.to_thread(self._func, inputs)


class LLMNode(Node):
    """Calls an LLM with a prompt template rendered from upstream inputs."""

    def __init__(
        self,
        name: str,
        llm: BackendAdapter,
        prompt_template: str,
        description: str = "",
    ) -> None:
        super().__init__(name, description)
        self._llm = llm
        self._template = prompt_template

    async def execute(self, inputs: dict[str, Any]) -> str:
        from thryve.core.context.models import Message

        prompt = self._template.format(**inputs)
        response = await self._llm.generate_async([Message(role="user", content=prompt)])
        return response.content


class AgentNode(Node):
    """Runs a full agent loop as a graph node."""

    def __init__(self, name: str, agent: Agent, description: str = "") -> None:
        super().__init__(name, description)
        self._agent = agent

    async def execute(self, inputs: dict[str, Any]) -> Any:
        task = inputs.get("task", str(inputs))
        result = await self._agent.run(task)
        return result.final_response


class ToolNode(Node):
    """Executes a single tool as a graph node."""

    def __init__(self, name: str, tool: Tool, description: str = "") -> None:
        super().__init__(name, description)
        self._tool = tool

    async def execute(self, inputs: dict[str, Any]) -> Any:
        import asyncio
        import inspect

        handler = self._tool.handler
        if handler is None:
            raise RuntimeError(f"Tool '{self._tool.name}' has no handler")
        if inspect.iscoroutinefunction(handler):
            return await handler(**inputs)
        return await asyncio.to_thread(handler, **inputs)
