from collections import deque
from typing import Any

from behave import given, then, when

from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
from cleveragents.langgraph.pure_graph import PureGraph


def _build_graph(name: str, node_defs: list[tuple[str, str | None]]):
    nodes: dict[str, NodeConfig] = {
        "start": NodeConfig(name="start", type=NodeType.START),
        "end": NodeConfig(name="end", type=NodeType.END),
    }
    edges: list[Edge] = []
    previous = "start"
    for node_name, function_name in node_defs:
        nodes[node_name] = NodeConfig(
            name=node_name, type=NodeType.FUNCTION, function=function_name
        )
        edges.append(Edge(source=previous, target=node_name))
        previous = node_name
    edges.append(Edge(source=previous, target="end"))
    return PureGraph(name=name, nodes=nodes, edges=edges)


@given('a pure graph with nodes "{first}" and "{second}"')
def step_graph_with_two_nodes(context, first: str, second: str):
    context.graph = _build_graph("pure-two", [(first, None), (second, None)])


@when("I list its topological order")
def step_list_topological_order(context):
    context.topo_order = context.graph.topological_order()


@then("the order should be start {first} {second} end")
def step_assert_topological_order(context, first: str, second: str):
    assert context.topo_order == ["start", first, second, "end"]


@given('a pure graph with function nodes "double" and "increment"')
def step_graph_with_functions(context):
    context.call_sequence = deque()

    def double(x: Any) -> Any:
        context.call_sequence.append("double")
        return x * 2

    def increment(x: Any) -> Any:
        context.call_sequence.append("increment")
        return x + 1

    context.fn_registry = {
        "double": double,
        "increment": increment,
    }
    context.graph = _build_graph(
        "pure-fns", [("double", "double"), ("increment", "increment")]
    )


@given("registered functions that double then increment")
def step_registered_functions(_context):
    # Functions already registered in previous step
    return


@when("I execute the pure graph starting with {initial:d}")
def step_execute_graph(context, initial: int):
    context.exec_result = context.graph.execute(context.fn_registry, initial=initial)


@then("the final result should be {expected:d}")
def step_assert_final_result(context, expected: int):
    assert context.exec_result == expected


@then("the functions should run in declaration order")
def step_assert_call_order(context):
    assert list(context.call_sequence) == ["double", "increment"]


@given("a pure graph with a missing function node")
def step_graph_with_missing_function(context):
    context.graph = _build_graph("pure-missing", [("skipper", "missing_fn")])
    context.fn_registry = {}


@then("the result should remain {expected:d}")
def step_assert_result_unchanged(context, expected: int):
    assert context.exec_result == expected


@given('a pure graph with non-functional nodes "{first}" and "{second}"')
def step_graph_with_non_functional_nodes(context, first: str, second: str):
    # Build graph without assigning functions and keep registry empty
    context.graph = _build_graph("pure-inert", [(first, None), (second, None)])
    context.fn_registry = {}