"""Shared utilities for pathway execution.

Operates on NodeInputs / NodeOutputs Pydantic models — no raw dicts.
"""

from typing import Any

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


def gather_inputs(
    pathway: Pathway,
    node_id: str,
    results: dict[str, NodeOutputs],
    initial_inputs: JsonObject,
) -> NodeInputs:
    """Gather inputs for a node from connections and initial inputs.

    Merges upstream NodeOutputs into a single NodeInputs model.
    Values are transferred *without* calling model_dump() so that typed
    objects (Pydantic models, typed lists, etc.) survive the trip through
    the graph intact.

    Args:
        pathway: The pathway being executed
        node_id: ID of the node to gather inputs for
        results: NodeOutputs from upstream nodes
        initial_inputs: Initial pathway inputs (dict from API boundary)

    Returns
    -------
        NodeInputs with merged upstream data
    """
    merged: dict[str, Any] = {}
    merged.update(initial_inputs)
    # Expose initial inputs under "input" namespace so templates
    # like {{ input.source_path }} resolve correctly.
    merged["input"] = dict(initial_inputs)

    for conn in pathway.connections:
        if conn.to_node == node_id:
            upstream = results.get(conn.from_node)
            if upstream is None:
                continue

            # Build a dict-like view of the upstream outputs *without*
            # serialisation, preserving typed Python objects.
            upstream_data: dict[str, Any] = dict(upstream.items())

            # Promote the "output" dict's fields so templates can use
            # {{node.response}} instead of {{node.output.response}}.
            # The inner keys take precedence over structural keys
            # (error, error_type, etc.) when there's a collision.
            output_val = upstream_data.get("output")
            if isinstance(output_val, dict):
                promoted = {**upstream_data, **output_val}
            else:
                promoted = upstream_data

            # Store under the node ID — supports both
            # {{search.response}} (promoted) and {{search.output.response}} (full).
            merged[conn.from_node] = promoted

            # Extract specific output field if requested
            value = upstream_data.get(conn.from_output, upstream_data)

            # Expose extracted value under to_input only when different
            # from from_node, so we don't overwrite the full result.
            target_key = conn.to_input if conn.to_input != "input" else conn.from_node
            if target_key != conn.from_node:
                merged[target_key] = value

    return NodeInputs(**merged)


def should_execute(
    pathway: Pathway,
    node_id: str,
    results: dict[str, NodeOutputs],
    skipped: set[str] | None = None,
) -> bool:
    """Check if node should execute (not gated out).

    Parameters
    ----------
    pathway
        The pathway being executed.
    node_id
        ID of the node to check.
    results
        NodeOutputs from upstream nodes.
    skipped
        Set of node IDs that were skipped (completed without executing).
        When **all** upstream connections come from skipped nodes, the
        current node is transitively skipped as well.

    Returns
    -------
    bool
        True if node should execute, False if gated out.
    """
    _skipped = skipped or set()

    upstream_connections = [conn for conn in pathway.connections if conn.to_node == node_id]

    # ConditionalNode sets ``selected_route`` in its output; if the upstream
    # node routed elsewhere, this node should be skipped.
    for conn in upstream_connections:
        upstream = results.get(conn.from_node)
        if (
            upstream is not None
            and "selected_route" in upstream
            and upstream["selected_route"] != node_id
        ):
            return False

    # Transitive skip propagation: if every upstream node was skipped
    # (present in `skipped`, absent from `results`), skip this node too.
    return not (
        upstream_connections and all(conn.from_node in _skipped for conn in upstream_connections)
    )
