import asyncio
from unittest.mock import AsyncMock, MagicMock

from behave import given, then, when
from rx.scheduler.eventloop import AsyncIOScheduler

from cleveragents.langgraph.graph import GraphConfig, LangGraph
from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
from cleveragents.langgraph.state import GraphState
from cleveragents.reactive.stream_router import StreamMessage

# Existing steps cover constructor scheduling, execution dispatch, start validation,
# cycle detection, and node executor state updates.


@given("a langgraph config with layered nodes")
def step_layered_config(context):
    nodes = {
        "a": NodeConfig(name="a", type=NodeType.FUNCTION),
        "b": NodeConfig(name="b", type=NodeType.FUNCTION),
        "c": NodeConfig(name="c", type=NodeType.FUNCTION),
    }
    edges = [Edge(source="a", target="b"), Edge(source="b", target="c")]
    context.graph = _fresh_graph(
        GraphConfig(name="graph-levels", nodes=nodes, edges=edges)
    )


@when("I compute the topological levels")
def step_compute_levels(context):
    context.levels = context.graph._topological_levels()  # pylint: disable=protected-access


@then("the levels should list each layer")
def step_assert_levels(context):
    assert context.levels[0] == {"a"}
    assert context.levels[1] == {"b"}
    assert context.levels[2] == {"c"}


@given("a langgraph instance with recorded history")
def step_history_instance(context):
    config = GraphConfig(name="graph-history")
    context.graph = _fresh_graph(config)
    context.graph.execution_history = ["x", "y"]


@when("I request execution history")
def step_get_history(context):
    context.history_result = context.graph.get_execution_history()


@then("it should return a copy of the recorded history")
def step_assert_history_copy(context):
    assert context.history_result == ["x", "y"]
    assert context.history_result is not context.graph.execution_history


@given("a langgraph config with explicit start and end nodes")
def step_explicit_guards(context):
    start = NodeConfig(name="start", type=NodeType.START, metadata={"guard": True})
    end = NodeConfig(name="end", type=NodeType.END, metadata={"guard": True})
    nodes = {
        "start": start,
        "end": end,
        "worker": NodeConfig(name="worker", type=NodeType.FUNCTION),
    }
    context.graph = _fresh_graph(GraphConfig(name="graph-guards", nodes=nodes))


@when("I construct the LangGraph with that config")
def step_construct_with_guards(_context):
    # Construction already performed in the given step
    return


@then("initialization should reuse the provided start and end nodes")
def step_assert_guard_reuse(context):
    assert context.graph.nodes["start"].config is context.graph.config.nodes["start"]
    assert context.graph.nodes["end"].config is context.graph.config.nodes["end"]


@given("a langgraph instance prepared to observe node streams")
def step_prepare_node_stream(context):
    nodes = {"worker": NodeConfig(name="worker", type=NodeType.FUNCTION)}
    context.graph = _fresh_graph(GraphConfig(name="graph-observe", nodes=nodes))
    context.worker_stream = context.graph.stream_router.streams[
        f"__{context.graph.name}_node_worker__"
    ]
    context.graph.logger = MagicMock()


@when("the node stream emits next error and completion events")
def step_emit_stream_events(context):
    context.worker_stream.on_next("payload")
    context.worker_stream.on_error(ValueError("boom"))
    context.worker_stream.on_completed()


@then("the subscription handlers should be exercised")
def step_assert_stream_handlers(context):
    context.graph.logger.error.assert_called_once()


@given("a langgraph config with converging branches")
def step_converging_config(context):
    nodes = {
        "a": NodeConfig(name="a", type=NodeType.FUNCTION),
        "b": NodeConfig(name="b", type=NodeType.FUNCTION),
        "c": NodeConfig(name="c", type=NodeType.FUNCTION),
    }
    edges = [
        Edge(source="start", target="a"),
        Edge(source="start", target="b"),
        Edge(source="a", target="c"),
        Edge(source="b", target="c"),
        Edge(source="c", target="end"),
    ]
    context.graph = _fresh_graph(
        GraphConfig(name="graph-parallel", nodes=nodes, edges=edges)
    )


@when("I construct the LangGraph for parallel groups")
def step_construct_parallel(_context):
    return


@then("parallel groups should reflect staged execution order")
def step_assert_parallel_groups(context):
    assert context.graph.parallel_groups[0] == {"start"}
    assert context.graph.parallel_groups[1] == {"a", "b"}
    assert context.graph.parallel_groups[2] == {"c"}
    assert context.graph.parallel_groups[3] == {"end"}


@given("a langgraph instance missing required guard nodes")
def step_missing_guard_nodes(context):
    config = GraphConfig(name="graph-missing-guards")
    context.graph = _fresh_graph(config)
    context.graph.nodes.pop("start", None)


@when("I validate the graph explicitly")
def step_validate_graph(context):
    context.validation_error = None
    try:
        context.graph._validate_graph()  # pylint: disable=protected-access
    except Exception as exc:  # pylint: disable=broad-except
        context.validation_error = exc


@then("a guard validation error should be raised")
def step_assert_guard_validation(context):
    assert isinstance(context.validation_error, ValueError)


@given("a langgraph instance ready to start")
def step_ready_to_start(context):
    config = GraphConfig(name="graph-start")
    context.graph = _fresh_graph(config)
    context.start_stream = f"__{context.graph.name}_node_start__"
    context.graph.stream_router.send_message = MagicMock()


@when("I start without input then stop and fetch state")
def step_start_and_stop(context):
    context.graph.start()
    context.graph.stop()
    context.final_state = context.graph.get_state()


@then("it should use default state and update running flags")
def step_assert_start_stop(context):
    context.graph.stream_router.send_message.assert_called_with(
        context.start_stream, GraphState()
    )
    assert context.graph.is_running is False
    assert isinstance(context.final_state, GraphState)


def _fresh_graph(config: GraphConfig) -> LangGraph:
    """Construct a LangGraph with isolated scheduler and router."""
    return LangGraph(config=config, stream_router=None, scheduler=None)


@given("a langgraph config with only start and end nodes")
def step_minimal_config(context):
    context.graph_config = GraphConfig(name="graph-minimal")


@given("asyncio has no running loop")
def step_no_running_loop(_context):
    try:
        loop = asyncio.get_running_loop()
        if loop.is_running():
            loop.stop()
    except RuntimeError:
        # Expected when no loop is running
        pass


@when("I construct a LangGraph without providing a scheduler")
def step_construct_graph(context):
    context.graph = _fresh_graph(context.graph_config)


@then("it should create an AsyncIOScheduler and initialize streams")
def step_assert_scheduler_and_streams(context):
    assert isinstance(context.graph.scheduler, AsyncIOScheduler)
    start_stream = f"__{context.graph.name}_node_start__"
    assert start_stream in context.graph.stream_router.streams


@given("a langgraph instance with a fresh state manager")
def step_graph_with_state_manager(context):
    config = GraphConfig(name="graph-exec")
    context.graph = _fresh_graph(config)
    context.start_stream = f"__{context.graph.name}_node_start__"
    context.graph.stream_router.send_message = MagicMock()


@when("I execute the graph with input data")
def step_execute_graph(context):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    result = loop.run_until_complete(context.graph.execute({"messages": []}))
    loop.close()
    asyncio.set_event_loop(asyncio.new_event_loop())
    context.execute_result = result


@then("the start stream should receive the state and execution returns a GraphState")
def step_assert_execute(context):
    context.graph.stream_router.send_message.assert_called_once()
    called_state = context.graph.stream_router.send_message.call_args.args[1]
    assert isinstance(called_state, GraphState)
    assert isinstance(context.execute_result, GraphState)


@given("a langgraph instance with a removed start stream")
def step_graph_missing_start(context):
    config = GraphConfig(name="graph-missing-start")
    context.graph = _fresh_graph(config)
    start_stream = f"__{context.graph.name}_node_start__"
    context.graph.stream_router.streams.pop(start_stream, None)


@when("I call start on the graph")
def step_call_start(context):
    context.start_error = None
    try:
        context.graph.start()
    except Exception as exc:  # pylint: disable=broad-except
        context.start_error = exc


@then("a start stream error should be raised")
def step_assert_start_error(context):
    assert isinstance(context.start_error, ValueError)


@given("a langgraph config with parallel execution disabled and a cycle")
def step_cycle_config(context):
    nodes = {
        "a": NodeConfig(name="a", type=NodeType.FUNCTION),
    }
    edges = [Edge(source="start", target="a"), Edge(source="a", target="start")]
    context.graph_config = GraphConfig(
        name="graph-cycle",
        nodes=nodes,
        edges=edges,
        parallel_execution=False,
    )


@when("I construct the LangGraph")
def step_construct_cycle_graph(context):
    context.graph = _fresh_graph(context.graph_config)


@then("the graph should report cycles and empty parallel groups")
def step_assert_cycles(context):
    assert context.graph.has_cycles is True
    assert context.graph.parallel_groups == []


@given("a langgraph instance with a patched state manager")
def step_graph_with_executor(context):
    nodes = {
        "worker": NodeConfig(name="worker", type=NodeType.FUNCTION),
    }
    config = GraphConfig(name="graph-executor", nodes=nodes)
    graph = _fresh_graph(config)
    graph.state_manager.get_state = MagicMock(return_value=GraphState())
    graph.state_manager.update_state = MagicMock()
    graph.nodes["worker"].execute = AsyncMock(return_value={"messages": ["m"]})
    context.executor = graph.stream_router._builtin_execute_node_worker
    context.graph = graph


@when("I invoke the builtin node executor")
def step_invoke_executor(context):
    result = context.executor(StreamMessage(content="x", metadata={}))
    context.executor_result = result


@then("it should update state and record execution history")
def step_assert_executor(context):
    context.graph.state_manager.update_state.assert_called_once()
    assert context.graph.execution_history == ["worker"]
    assert isinstance(context.executor_result, StreamMessage)
