"""Node parsing — YAML dict to executable node, no intermediate spec classes.

Usage:
    from core.pathways.models.node_spec import parse_node_spec

    # Returns an executable node directly (LLMNode, ToolNode, etc.)
    node = parse_node_spec({"id": "search", "type": "tool", "tool": "web.search"})
"""

from __future__ import annotations

from typing import Any

from core.pathways.nodes import (
    CheckpointNode,
    CodeExecuteNode,
    LLMNode,
    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.nodes.handoff import HandoffCandidate, HandoffNode

# Map from YAML ``type`` string to the executable node class.
_NODE_TYPE_MAP: dict[str, type[NodeBase]] = {
    "llm": LLMNode,
    "tool": ToolNode,
    "transform": TransformNode,
    "input": InputNode,
    "output": OutputNode,
    "pathway": PathwayNode,
    "checkpoint": CheckpointNode,
    "code_execute": CodeExecuteNode,
    "conditional": ConditionalNode,
    "handoff": HandoffNode,
    "stage": StageNode,
}


def _normalize_yaml_dict(data: dict[str, Any]) -> dict[str, Any]:
    """Flatten a YAML node dict for direct Pydantic construction.

    YAML authors may use either:
    - Flat form:   ``{id: x, type: llm, prompt: "hello", model: fast}``
    - Config form: ``{id: x, type: llm, config: {prompt: "hello", model: fast}}``

    This function merges ``config`` sub-dict into the top level so that
    node Pydantic models (which declare fields at top level) can parse it.
    """
    out = dict(data)
    config = out.pop("config", None)
    if isinstance(config, dict):
        for k, v in config.items():
            out.setdefault(k, v)

    # Normalise common aliases
    if "tool_name" in out and "tool" not in out:
        out["tool"] = out.pop("tool_name")
    if "expression" in out and "expr" not in out:
        out["expr"] = out.pop("expression")
    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")

    # HandoffNode: pre-parse candidate dicts
    if out.get("type") == "handoff" and "candidates" in out:
        out["candidates"] = [
            HandoffCandidate.model_validate(c) if isinstance(c, dict) else c
            for c in out["candidates"]
        ]

    # StageNode: normalise transitions from list-of-lists to list-of-tuples
    if out.get("type") == "stage" and "transitions" in out:
        raw = out["transitions"]
        out["transitions"] = [
            (str(t[0]), str(t[1]))
            for t in raw
            if isinstance(t, (list, tuple)) and len(t) == 2
        ]

    return out


def parse_node_spec(data: dict[str, Any]) -> NodeBase:
    """Parse a YAML dict into an executable node.

    Parameters
    ----------
    data
        Dict from YAML with at least a ``type`` field (defaults to ``"llm"``).

    Returns
    -------
    NodeBase
        Ready-to-execute node instance.

    Raises
    ------
    ValueError
        If the ``type`` is unknown or Pydantic validation fails.
    """
    node_type = data.get("type", "llm")
    node_class = _NODE_TYPE_MAP.get(node_type)
    if node_class is None:
        msg = f"Unknown node type: {node_type}"
        raise ValueError(msg)

    flat = _normalize_yaml_dict(data)
    return node_class.model_validate(flat)


# Keep the old names alive so existing imports keep working.
NodeSpec = NodeBase
NodeSpecBase = NodeBase


__all__ = [
    "NodeSpec",
    "NodeSpecBase",
    "parse_node_spec",
]
