"""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 collections.abc import Iterator
from typing import TYPE_CHECKING
from uuid import uuid4

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

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

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from core.pathways.context import Context
    from core.pathways.pathway import Pathway


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  # noqa: PLC0415 (circular)

        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  # noqa: PLC0415 (circular)

        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  # noqa: PLC0415 (circular)

        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_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",
]
