"""Pathway - The execution graph structure.

This module defines:
- Pathway: A graph of nodes and connections (serializable!)
- Connection: Links between nodes
- Signal: Result of node execution

Control flow is handled by node types (ConditionalNode, LoopNode).
Legacy Gate/Loop models have been removed — use ConditionalNode and
LoopNode via PathwayBuilder.conditional() and .loop_node() instead.

Nodes are defined in nodes/core.py (LLMNode, ToolNode, etc.)

Serialization:
    # Pathway to JSON
    pathway = Pathway(id="example", name="Example")
    json_data = pathway.model_dump_json()

    # JSON to Pathway
    pathway2 = Pathway.model_validate_json(json_data)

    # Execute
    await vm.execute(pathway2, inputs)
"""

import importlib
import logging
from enum import Enum, StrEnum
from typing import Annotated, Any, Self
from uuid import uuid4

from pydantic import BaseModel, Field, model_validator

# =============================================================================
# HOISTED IMPORTS - All imports at module level (Rule 1: No lazy imports)
# =============================================================================
from core.pathways.expressions import safe_eval

logger = logging.getLogger(__name__)

# Import node types for discriminated union serialization
# This must be a runtime import (not TYPE_CHECKING) for Pydantic to work
#
# NODE TYPES:
# Foundational:
# - InputNode: Declares pathway inputs with schema validation
# - OutputNode: Declares pathway outputs with mapping
#
# Processing:
# - ToolNode: Tool invocation node (calls Python-backed tools)
# - LLMNode: Core compute primitive for direct LLM generation
# - TransformNode: Safe expression evaluation
#
# Composition:
# - PathwayNode: Execute a sub-pathway (pathways in pathways)
#
# Stateful (for orchestration):
# - StageNode: Orchestration state with agent + transitions
# - LoopNode: Iteration until condition
# - CheckpointNode: Persistence point for resume
#
from core.common import JsonObject, JSONValue
from core.pathways.nodes import (
    CheckpointNode,
    LLMNode,
    LoopNode,
    PathwayNode,
    StageNode,
    ToolNode,
    TransformNode,
)
from core.pathways.nodes.base import NodeBase
from core.pathways.nodes.flow import ConditionalNode
from core.pathways.nodes.foundational import InputNode, OutputNode
from core.pathways.schemas.signature import PathwaySignature

# Pseudo-node IDs that can be used in connections without explicit node objects.
# If a pathway has an explicit InputNode/OutputNode with these IDs, they're real nodes.
# If not, the VM handles them as implicit entry/exit points.
PSEUDO_NODE_IDS: frozenset[str] = frozenset({"input", "output"})


def is_real_node(node_id: str, pathway_nodes: "dict[str, Node] | None" = None) -> bool:
    """Check if a node ID refers to a real executable node.

    Args:
        node_id: The node ID to check
        pathway_nodes: Optional dict of pathway nodes for additional validation

    Returns
    -------
        True if node_id is a real node (either explicitly added or not a pseudo-node ID)
    """
    # If pathway_nodes is provided and contains this ID, it's a real node
    if pathway_nodes is not None:
        if node_id in pathway_nodes:
            return True
        # Not in pathway_nodes - check if it's a pseudo-node ID
        return node_id not in PSEUDO_NODE_IDS

    # Without pathway context, assume pseudo-node IDs are not real
    return node_id not in PSEUDO_NODE_IDS


# Discriminated union of node types - enables JSON serialization/deserialization
Node = Annotated[
    InputNode
    | OutputNode
    | ToolNode
    | TransformNode
    | LLMNode
    | PathwayNode
    | StageNode
    | LoopNode
    | CheckpointNode
    | ConditionalNode,
    Field(discriminator="type"),
]


class NodeStatus(Enum):
    """Status of a node execution."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"


class Signal(BaseModel):
    """Result of a node execution."""

    model_config = {"extra": "allow"}

    value: JSONValue = None
    error: str | None = None
    status: NodeStatus = NodeStatus.PENDING
    metadata: JsonObject = Field(default_factory=dict)

    @property
    def success(self) -> bool:
        return self.status == NodeStatus.COMPLETED and self.error is None


# =============================================================================
# INPUT NORMALIZATION HELPERS
# =============================================================================

_NODE_ALIASES: dict[str, str] = {
    "tool_name": "tool",
    "expression": "expr",
}


def _normalize_node_dict(node_id: str, data: dict[str, Any]) -> dict[str, Any]:
    """Normalize a single node dict for direct Pydantic construction.

    Handles:
    - Injecting ``id`` from the dict key
    - Flattening ``config`` sub-dict into top level
    - Common field aliases (``tool_name`` -> ``tool``, ``expression`` -> ``expr``)
    - ``schema`` alias for InputNode/OutputNode
    """
    out = dict(data)
    out.setdefault("id", node_id)

    # Flatten config sub-dict
    config = out.pop("config", None)
    if isinstance(config, dict):
        for k, v in config.items():
            out.setdefault(k, v)

    # Normalise common aliases
    for old_key, new_key in _NODE_ALIASES.items():
        if old_key in out and new_key not in out:
            out[new_key] = out.pop(old_key)

    # schema → input_schema / output_schema depending on node type
    if "schema" in out:
        node_type = out.get("type", "llm")
        if node_type == "input" and "input_schema" not in out:
            out["input_schema"] = out.pop("schema")
        elif node_type == "output" and "output_schema" not in out:
            out["output_schema"] = out.pop("schema")

    return out


def _parse_connection_string(s: str) -> dict[str, str] | None:
    """Parse ``"a -> b"`` or ``"a → b"`` into a connection dict."""
    if "→" in s:
        parts = s.split("→")
    elif "->" in s:
        parts = s.split("->")
    else:
        return None
    if len(parts) == 2:
        return {"from_node": parts[0].strip(), "to_node": parts[1].strip()}
    return None


# =============================================================================
# CONNECTION
# =============================================================================


class Connection(BaseModel):
    """Connection between nodes.

    For stateful pathways, connections can have conditional transitions:
    - when: Expression that must evaluate to true for this connection to be followed

    Accepts both YAML-style (``from``/``to``) and Python-style
    (``from_node``/``to_node``) keys.

    Example:
        Connection(from_node="review", to_node="approve", when="{{review.passed}}")
        Connection(from_node="review", to_node="revise", when="{{not review.passed}}")
    """

    model_config = {"extra": "allow"}

    from_node: str
    to_node: str

    @model_validator(mode="before")
    @classmethod
    def _normalize_aliases(cls, data: Any) -> Any:
        """Accept ``from``/``to`` as aliases for ``from_node``/``to_node``."""
        if not isinstance(data, dict):
            return data
        if "from" in data and "from_node" not in data:
            data["from_node"] = data.pop("from")
        if "to" in data and "to_node" not in data:
            data["to_node"] = data.pop("to")
        return data

    from_output: str = "output"
    to_input: str = "input"

    # Conditional transition (for stateful pathways)
    # If set, this connection is only followed when the expression evaluates to true
    when: str | None = None

    def transfer(self, data: JSONValue) -> JSONValue:
        """Transfer data through this connection."""
        return data

    def should_follow(self, context: dict[str, JSONValue]) -> bool:
        """Check if this connection should be followed based on the 'when' condition.

        Args:
            context: Current execution context with node results

        Returns
        -------
            True if connection should be followed (no condition or condition is true)
        """
        if self.when is None:
            return True

        # Using hoisted import: safe_eval
        # DSL Unification: 'state' is an alias for 'context'
        enriched_context = dict(context)
        if "context" in enriched_context and "state" not in enriched_context:
            enriched_context["state"] = enriched_context["context"]
        elif "state" in enriched_context and "context" not in enriched_context:
            enriched_context["context"] = enriched_context["state"]

        try:
            return bool(safe_eval(self.when, enriched_context))
        except Exception:
            logger.debug(
                "Connection.should_follow: condition eval failed for '%s'",
                self.when,
                exc_info=True,
            )
            return False


class PathwayMode(StrEnum):
    """Pathway execution mode.

    - dag: Directed Acyclic Graph - parallel execution, stateless, no cycles (default)
    - stateful: Sequential, persistent state, cycles allowed, multi-turn capable
    """

    DAG = "dag"
    STATEFUL = "stateful"


class Pathway(BaseModel):
    """A computation graph - fully serializable.

    THE UNIVERSAL COMPOSITION LANGUAGE of AlbusOS.

    A Pathway is:
    - nodes: Dict of node_id -> Node (discriminated union, serializable)
    - connections: List of Connection (data flow between nodes)
    - mode: Execution mode (dag or stateful)
    - state_schema: Optional Pydantic model path for type-safe state validation

    Control flow (conditionals, loops) is expressed as nodes
    (``ConditionalNode``, ``LoopNode``) — not as separate gate/loop lists.

    Modes:
    - dag (default): Parallel execution, stateless, no cycles
    - stateful: Sequential, persistent, cycles allowed, multi-turn

    Composition (DSL):
    - pathway >> node: append node to end
    - pathway >> other: chain pathways
    - pathway | other: parallel composition

    Serialization:
    - pathway.model_dump_json() → JSON string
    - Pathway.model_validate_json(json) → Pathway object

    Execution happens in PathwayVM.
    """

    model_config = {"extra": "allow"}

    @model_validator(mode="before")
    @classmethod
    def _normalize_input(cls, data: Any) -> Any:
        """Normalize raw dicts so Pydantic can validate them directly.

        Handles:
        - Injecting node ``id`` from dict keys
        - Flattening ``config`` sub-dicts
        - Field aliases (``tool_name``, ``expression``, ``schema``)
        - String-format connections (``"a -> b"``)
        - List-format nodes (``[{id: ..., type: ...}]``)
        """
        if not isinstance(data, dict):
            return data

        # --- Normalize nodes ---
        raw_nodes = data.get("nodes")
        if isinstance(raw_nodes, dict):
            normalized: dict[str, Any] = {}
            for node_id, node_data in raw_nodes.items():
                normalized_data = (
                    _normalize_node_dict(node_id, node_data)
                    if isinstance(node_data, dict)
                    else node_data
                )
                normalized[node_id] = normalized_data
            data["nodes"] = normalized
        elif isinstance(raw_nodes, list):
            # List format: [{id: ..., type: ..., ...}]
            normalized = {}
            for item in raw_nodes:
                if isinstance(item, dict):
                    nid = item.get("id", f"n_{uuid4().hex[:6]}")
                    normalized[nid] = _normalize_node_dict(nid, item)
            data["nodes"] = normalized

        # --- Normalize connections ---
        raw_conns = data.get("connections")
        if isinstance(raw_conns, list):
            normalized_conns: list[Any] = []
            for c in raw_conns:
                if isinstance(c, str):
                    parsed = _parse_connection_string(c)
                    if parsed:
                        normalized_conns.append(parsed)
                else:
                    normalized_conns.append(c)
            data["connections"] = normalized_conns

        return data

    id: str = Field(default_factory=lambda: f"pathway_{uuid4().hex[:8]}")
    name: str | None = None
    description: str | None = None
    version: str | None = None  # Semantic version for the pathway

    # Execution mode - determines how the pathway runs
    mode: PathwayMode = PathwayMode.DAG

    # State schema - path to Pydantic model for type-safe state validation.
    # When set, initial context is validated against this schema.
    state_schema: str | None = None

    # Discriminated union enables proper JSON serialization/deserialization
    nodes: dict[str, Node] = Field(default_factory=dict)

    connections: list[Connection] = Field(default_factory=list)
    metadata: JsonObject = Field(default_factory=dict)

    # Typed I/O contract (optional)
    signature: PathwaySignature | None = None

    # Trigger configuration (optional) - when to auto-run this pathway.
    # Accepts a TriggerSpec dict (type + config) or raw JsonObject for
    # backward compatibility.
    trigger: JsonObject | None = None

    # =========================================================================
    # STATE SCHEMA VALIDATION
    # =========================================================================

    def validate_state(self, state: JsonObject) -> tuple[bool, list[str]]:
        """Validate state against the state_schema if defined.

        This enables type-safe state management for complex pathways.
        The schema can be a Pydantic model path like "core.models.TriageState".

        Args:
            state: The state dict to validate

        Returns
        -------
            Tuple of (is_valid, list_of_errors)
        """
        if not self.state_schema:
            return True, []

        try:
            # Dynamically import the schema class
            parts = self.state_schema.rsplit(".", 1)
            if len(parts) != 2:
                return False, [
                    f"Invalid state_schema format: {self.state_schema}. Expected 'module.ClassName'",
                ]

            module_path, class_name = parts

            # Using hoisted import: importlib
            module = importlib.import_module(module_path)
            schema_class = getattr(module, class_name, None)

            if schema_class is None:
                return False, [f"State schema class not found: {self.state_schema}"]

            # Validate using Pydantic (BaseModel already imported at module level)
            if not (isinstance(schema_class, type) and issubclass(schema_class, BaseModel)):
                return False, [f"State schema must be a Pydantic BaseModel: {self.state_schema}"]

            # Attempt validation
            schema_class.model_validate(state)
            return True, []

        except ImportError as e:
            return False, [f"Failed to import state_schema module: {e}"]
        except Exception as e:
            # Pydantic validation errors
            error_msg = str(e)
            return False, [f"State validation failed: {error_msg}"]

    def get_state_schema_class(self) -> type | None:
        """Get the Pydantic class for state validation if defined.

        Returns
        -------
            The Pydantic model class, or None if no schema is defined
        """
        if not self.state_schema:
            return None

        try:
            parts = self.state_schema.rsplit(".", 1)
            if len(parts) != 2:
                return None

            module_path, class_name = parts

            # Using hoisted import: importlib
            module = importlib.import_module(module_path)
            return getattr(module, class_name, None)
        except Exception:
            return None

    # =========================================================================
    # ENTRY/EXIT for composition
    # =========================================================================

    @property
    def entry_id(self) -> str:
        """ID of the entry node (first node with no incoming connections)."""
        if not self.nodes:
            raise ValueError("Pathway has no nodes")

        has_incoming = {conn.to_node for conn in self.connections}
        entry_candidates = [nid for nid in self.nodes if nid not in has_incoming]

        if entry_candidates:
            return entry_candidates[0]
        return next(iter(self.nodes.keys()))

    @property
    def exit_id(self) -> str:
        """ID of the exit node (last node with no outgoing connections)."""
        if not self.nodes:
            raise ValueError("Pathway has no nodes")

        has_outgoing = {conn.from_node for conn in self.connections}
        exit_candidates = [nid for nid in self.nodes if nid not in has_outgoing]

        if exit_candidates:
            return exit_candidates[-1]
        return list(self.nodes.keys())[-1]

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

    def __rshift__(self, other: "NodeBase | Pathway") -> Self:
        """Chain: self >> other."""
        if isinstance(other, Pathway):
            return Pathway(
                id=f"chain_{uuid4().hex[:6]}",
                nodes={**self.nodes, **other.nodes},
                connections=[
                    *self.connections,
                    Connection(from_node=self.exit_id, to_node=other.entry_id),
                    *other.connections,
                ],
                metadata={**self.metadata, **other.metadata},
            )

        # It's a node
        return Pathway(
            id=f"chain_{uuid4().hex[:6]}",
            nodes={**self.nodes, other.id: other},
            connections=[
                *self.connections,
                Connection(from_node=self.exit_id, to_node=other.id),
            ],
            metadata=dict(self.metadata),
        )

    def __or__(self, other: "Pathway") -> Self:
        """Parallel: self | other."""
        if not isinstance(other, Pathway):
            raise TypeError(f"Cannot parallel compose with {type(other)}")

        parallel_group = self.metadata.get("parallel_group", list(self.nodes.keys()))
        other_group = other.metadata.get("parallel_group", list(other.nodes.keys()))

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

    # =========================================================================
    # BUILDER METHODS
    # =========================================================================

    def add_node(self, node_id_or_node: str | NodeBase, node: NodeBase | None = None) -> Self:
        """Add a node to this pathway.

        Supports two calling patterns:
            # Pattern 1: ID + node (recommended)
            pathway.add_node("input", InputNode(input_schema={...}))

            # Pattern 2: node with id field
            pathway.add_node(InputNode(id="input", input_schema={...}))
        """
        if isinstance(node_id_or_node, str):
            # Pattern 1: add_node("id", node)
            if node is None:
                raise ValueError(
                    "When first argument is a string (node_id), second argument (node) is required",
                )
            node_id = node_id_or_node
            node.id = node_id  # type: ignore[union-attr]
            self.nodes[node_id] = node  # type: ignore[assignment]
        else:
            # Pattern 2: add_node(node) where node has an id
            the_node = node_id_or_node
            self.nodes[the_node.id] = the_node  # type: ignore[assignment]
        return self

    def connect(self, from_node: str, to_node: str) -> Self:
        """Connect two nodes."""
        self.connections.append(Connection(from_node=from_node, to_node=to_node))
        return self


__all__ = [
    # Constants
    "PSEUDO_NODE_IDS",
    # Types
    "Connection",
    "Node",
    "NodeStatus",
    "Pathway",
    "PathwayMode",
    "Signal",
    # Functions
    "is_real_node",
]
