"""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())

            # Store full upstream result under the node ID
            # so templates like {{search.output}} work.
            merged[conn.from_node] = upstream_data

            # 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],
) -> bool:
    """Check if node should execute (not gated out).

    Args:
        pathway: The pathway being executed
        node_id: ID of the node to check
        results: NodeOutputs from upstream nodes

    Returns
    -------
        True if node should execute, False if gated out
    """
    for gate in pathway.gates:
        if node_id in (gate.true_path, gate.false_path):
            for result in results.values():
                if "selected_route" in result:
                    selected = result["selected_route"]
                    if selected != node_id:
                        return False

    for conn in pathway.connections:
        if conn.to_node == node_id:
            upstream = results.get(conn.from_node)
            if (
                upstream is not None
                and "selected_route" in upstream
                and upstream["selected_route"] != node_id
            ):
                return False

    return True
