"""NodeBase - Base class for all node types with composition operators.

Nodes are composable values. They:
1. Are Pydantic models (serializable to JSON)
2. Have a compute() method that takes/returns Pydantic models
3. Support operators for composition (>>, |, &)

compute() signature:
    async def compute(self, inputs: NodeInputs, ctx: Context) -> NodeOutputs

- inputs: NodeInputs (Pydantic model with extra="allow") — gathered from upstream
- ctx: Context — execution context with tools, memory, services
- returns: NodeOutputs subclass — typed output model

Example:
    understand = LLMNode(prompt="Analyze: {{input}}")
    search = ToolNode(tool="web.search")

    # Chain with >>
    graph = understand >> search

    # Parallel with |
    multi = search_a | search_b | search_c
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any
from uuid import uuid4

from jinja2 import BaseLoader, Environment, Undefined
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from collections.abc import Iterator

    from core.common import JsonObject, JSONValue
    from core.pathways.models.io import NodeInputs, NodeOutputs
    from core.pathways.pathway import Pathway

    from .context import Context

logger = logging.getLogger(__name__)


class NodeBase(BaseModel):
    """Base class for all node types.

    Subclasses must:
    1. Define a `type` literal field (for serialization)
    2. Implement `compute(inputs, ctx)` method

    Example:
        class LLMNode(NodeBase):
            type: Literal["llm"] = "llm"
            prompt: str
            model: str = "reasoning"  # Capability name or explicit model

            async def compute(self, inputs: NodeInputs, ctx: Context) -> NodeOutputs:
                ...
    """

    model_config = {"extra": "forbid"}

    # Identity
    id: str = Field(default_factory=lambda: f"n_{uuid4().hex[:6]}")
    name: str | None = None
    description: str | None = None

    # Execution resilience — overridable per node
    timeout_ms: int | None = Field(
        default=None,
        description="Per-node timeout in milliseconds. None = use VM default.",
    )
    retry_count: int = Field(
        default=0,
        ge=0,
        description="Number of retries on failure. 0 = no retries.",
    )
    retry_delay_ms: int = Field(
        default=1000,
        ge=0,
        description="Base delay between retries in ms (doubles each retry).",
    )

    # Subclasses must override
    async def compute(self, inputs: NodeInputs, ctx: Context) -> NodeOutputs:
        """Execute this node. Subclasses must implement.

        Args:
            inputs: NodeInputs model from upstream nodes (gathered by VM from connections)
            ctx: Execution context with tools, memory, services

        Returns
        -------
            NodeOutputs subclass. Convention: response, output, or domain-specific fields.
        """
        raise NotImplementedError(f"{self.__class__.__name__} must implement compute()")

    # =========================================================================
    # COMPOSITION OPERATORS
    # =========================================================================

    def __rshift__(self, other: NodeBase | Pathway) -> Pathway:
        """Chain: self >> other.

        Creates a Pathway where self connects to other.
        """
        from core.pathways.pathway import Connection, Pathway

        if isinstance(other, Pathway):
            # self >> existing_pathway: prepend self
            return Pathway(
                id=f"chain_{uuid4().hex[:6]}",
                nodes={self.id: self, **other.nodes},
                connections=[
                    Connection(from_node=self.id, to_node=other.entry_id),
                    *other.connections,
                ],
            )
        # self >> other_node: create two-node pathway
        return Pathway(
            id=f"chain_{uuid4().hex[:6]}",
            nodes={self.id: self, other.id: other},
            connections=[Connection(from_node=self.id, to_node=other.id)],
        )

    def __or__(self, other: NodeBase) -> Pathway:
        """Parallel: self | other.

        Creates a Pathway where both nodes run in parallel (no connection between them).
        They will receive the same input and their outputs can be joined downstream.
        """
        from core.pathways.pathway import Pathway

        return Pathway(
            id=f"parallel_{uuid4().hex[:6]}",
            nodes={self.id: self, other.id: other},
            connections=[],  # No connection = parallel
            metadata={"parallel_group": [self.id, other.id]},
        )

    def __and__(self, other: NodeBase) -> Pathway:
        """Join: self & other.

        Creates a Pathway indicating both must complete before continuing.
        Similar to parallel but semantic difference for scheduling.
        """
        from core.pathways.pathway import Pathway

        return Pathway(
            id=f"join_{uuid4().hex[:6]}",
            nodes={self.id: self, other.id: other},
            connections=[],
            metadata={"join_group": [self.id, other.id]},
        )

    # =========================================================================
    # TEMPLATE HELPERS  (Jinja2 — no regex)
    # =========================================================================

    # Shared Jinja2 environment: silent on missing vars, no auto-escaping.
    class _SilentUndefined(Undefined):
        """Return empty string for missing variables instead of raising."""

        def __str__(self) -> str:
            return ""

        def __iter__(self) -> Iterator[object]:
            return iter([])

        def __bool__(self) -> bool:
            return False

    _jinja_env = Environment(
        loader=BaseLoader(),
        undefined=_SilentUndefined,
        keep_trailing_newline=True,
    )

    def _format_template(self, template: str, data: JsonObject) -> str:
        """Render a Jinja2 template string against *data*.

        Supports the full Jinja2 syntax:
            {{ message }}                      — simple variable
            {{ node.output.response }}         — nested path
            {{ input.collection | default('x') }} — default filter
            {% if input.question %}…{% endif %} — conditionals
            {{ state.xxx }} / {{ context.xxx }} — aliases
        """
        # DSL unification: state ↔ context aliases
        enriched: dict[str, object] = dict(data)
        if "context" in enriched and "state" not in enriched:
            enriched["state"] = enriched["context"]
        elif "state" in enriched and "context" not in enriched:
            enriched["context"] = enriched["state"]

        try:
            tpl = self._jinja_env.from_string(template)
            return tpl.render(**enriched)
        except Exception:
            # Template syntax errors should not crash the pipeline —
            # return the raw template so the LLM still gets *something*.
            logger.debug(
                "Template render failed for node %s: %s",
                self.id,
                template[:80],
                exc_info=True,
            )
            return template

    def _resolve_path(self, data: JsonObject, path: str) -> Any:
        """Traverse a dot-path through a nested dict and return the value.

        Parameters
        ----------
        data
            The root dict to traverse (typically template context from inputs).
        path
            Dot-separated path, e.g. ``"search.output.tool_list"``.

        Returns
        -------
        Any
            The value at the path, or ``None`` if any segment is missing.
        """
        from functools import reduce

        path = path.strip()
        if not path:
            return None
        try:
            return reduce(
                lambda d, key: d[key] if isinstance(d, dict) else getattr(d, key),
                path.split("."),
                data,
            )
        except (KeyError, AttributeError, TypeError):
            logger.debug(
                "_resolve_path: path %r not found in data keys %s", path, list(data.keys())
            )
            return None

    def _resolve_templates(self, value: JSONValue, data: JsonObject) -> JSONValue:
        """Recursively resolve Jinja2 templates in nested structures.

        Args:
            value: Value to resolve (can be dict, list, str, or scalar)
            data: Context data for template resolution

        Returns
        -------
            Resolved value with all templates expanded
        """
        if isinstance(value, dict):
            return {k: self._resolve_templates(v, data) for k, v in value.items()}
        if isinstance(value, list):
            return [self._resolve_templates(item, data) for item in value]
        if isinstance(value, str) and ("{{" in value or "{%" in value):
            return self._format_template(value, data)
        return value


__all__ = [
    "NodeBase",
]
