"""Tests for Topology Integration - Categories 3-8 from Milestone 4 Test Plan."""
import asyncio
import pytest
from byzpy.engine.graph.pool import ActorPoolConfig


@pytest.fixture(autouse=True)
def clear_process_context_registry():
    """Clear ProcessContext registry before and after each test."""
    from byzpy.engine.node.context import ProcessContext
    ProcessContext._registry.clear()
    yield
    ProcessContext._registry.clear()


def make_app(name: str):
    """Create a simple NodeApplication for testing."""
    from byzpy.engine.node.application import NodeApplication
    return NodeApplication(
        name=name,
        actor_pool=[ActorPoolConfig(backend="thread", count=1)],
    )


# =============================================================================
# Category 3: DecentralizedNode Integration (6 tests)
# =============================================================================

# 3.1 Node Uses MessageRouter


@pytest.mark.asyncio
async def test_decentralizednode_uses_messagerouter():
    """Verify DecentralizedNode uses MessageRouter for message routing."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.decentralized import DecentralizedNode
    from byzpy.engine.node.router import MessageRouter
    from byzpy.engine.node.context import InProcessContext

    topology = Topology.ring(3, k=1)

    ctx = InProcessContext()
    app = make_app("test-app")

    node = DecentralizedNode(
        node_id=0,
        application=app,
        context=ctx,
        topology=topology,
    )

    assert isinstance(node.message_router, MessageRouter)
    assert node.message_router.topology is topology


@pytest.mark.asyncio
async def test_decentralizednode_send_to_neighbor():
    """Verify send_message() works for valid neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.decentralized import DecentralizedNode
    from byzpy.engine.node.context import InProcessContext

    topology = Topology.ring(3, k=1)

    contexts = [InProcessContext() for _ in range(3)]
    apps = [make_app(f"app-{i}") for i in range(3)]

    nodes = [
        DecentralizedNode(node_id=i, application=apps[i], context=contexts[i], topology=topology)
        for i in range(3)
    ]

    for node in nodes:
        await node.start()

    # Node 0 sends to neighbor 1
    await nodes[0].send_message(
        to_node_id=1,
        message_type="gradient",
        payload={"data": [1.0]},
    )

    # Give time for delivery
    await asyncio.sleep(0.1)

    # Verify node 1 received (via message handler or context)

    for node in nodes:
        await node.shutdown()


@pytest.mark.asyncio
async def test_decentralizednode_send_to_non_neighbor_raises():
    """Verify send_message() raises for non-neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.decentralized import DecentralizedNode
    from byzpy.engine.node.context import InProcessContext

    topology = Topology.ring(5, k=1)  # Node 0 can only send to 1 and 4

    ctx = InProcessContext()
    app = make_app("test-app")

    node = DecentralizedNode(
        node_id=0,
        application=app,
        context=ctx,
        topology=topology,
    )

    await node.start()

    # Try to send to node 2 (not a neighbor)
    with pytest.raises(ValueError, match="not a neighbor"):
        await node.send_message(to_node_id=2, message_type="test", payload={})

    await node.shutdown()


# 3.2 Broadcast and Multicast via Node


@pytest.mark.asyncio
async def test_decentralizednode_broadcast_message():
    """Verify broadcast_message() sends to all neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.decentralized import DecentralizedNode
    from byzpy.engine.node.context import InProcessContext

    topology = Topology.ring(4, k=1)

    contexts = [InProcessContext() for _ in range(4)]
    apps = [make_app(f"app-{i}") for i in range(4)]

    nodes = [
        DecentralizedNode(node_id=i, application=apps[i], context=contexts[i], topology=topology)
        for i in range(4)
    ]

    for node in nodes:
        await node.start()

    # Track received messages
    received_by = {i: [] for i in range(4)}
    for i, node in enumerate(nodes):
        async def handler(from_id, payload, idx=i):
            received_by[idx].append((from_id, payload))
        node.register_message_handler("broadcast_test", handler)

    # Node 0 broadcasts (neighbors are 1 and 3)
    await nodes[0].broadcast_message(
        message_type="broadcast_test",
        payload={"round": 1},
    )

    await asyncio.sleep(0.2)

    # Neighbors 1 and 3 should receive
    assert len(received_by[1]) == 1
    assert len(received_by[3]) == 1
    # Non-neighbor 2 should NOT receive
    assert len(received_by[2]) == 0

    for node in nodes:
        await node.shutdown()


@pytest.mark.asyncio
async def test_decentralizednode_multicast_message():
    """Verify multicast_message() sends to specified neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.decentralized import DecentralizedNode
    from byzpy.engine.node.context import InProcessContext

    topology = Topology.complete(4)  # Everyone is a neighbor

    contexts = [InProcessContext() for _ in range(4)]
    apps = [make_app(f"app-{i}") for i in range(4)]

    nodes = [
        DecentralizedNode(node_id=i, application=apps[i], context=contexts[i], topology=topology)
        for i in range(4)
    ]

    for node in nodes:
        await node.start()

    received_by = {i: [] for i in range(4)}
    for i, node in enumerate(nodes):
        async def handler(from_id, payload, idx=i):
            received_by[idx].append(payload)
        node.register_message_handler("multicast_test", handler)

    # Node 0 multicasts to only nodes 1 and 2
    await nodes[0].multicast_message(
        to_node_ids=[1, 2],
        message_type="multicast_test",
        payload={"partial": True},
    )

    await asyncio.sleep(0.2)

    assert len(received_by[1]) == 1
    assert len(received_by[2]) == 1
    assert len(received_by[3]) == 0  # Not in multicast list

    for node in nodes:
        await node.shutdown()


# 3.3 Get Neighbors API


@pytest.mark.asyncio
async def test_decentralizednode_get_out_neighbors():
    """Verify get_neighbors() returns node's outgoing neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.decentralized import DecentralizedNode
    from byzpy.engine.node.context import InProcessContext

    topology = Topology.ring(5, k=1)

    ctx = InProcessContext()
    app = make_app("test-app")

    node = DecentralizedNode(
        node_id=0,
        application=app,
        context=ctx,
        topology=topology,
    )

    # Node 0 in ring(5, k=1) has neighbors 1 and 4
    neighbors = node.get_neighbors()

    assert set(neighbors) == {1, 4}


@pytest.mark.asyncio
async def test_decentralizednode_get_in_neighbors():
    """Verify get_in_neighbors() returns nodes that can send to this node."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.decentralized import DecentralizedNode
    from byzpy.engine.node.context import InProcessContext

    topology = Topology.ring(5, k=1)

    ctx = InProcessContext()
    app = make_app("test-app")

    node = DecentralizedNode(
        node_id=0,
        application=app,
        context=ctx,
        topology=topology,
    )

    # In bidirectional ring, in-neighbors same as out-neighbors
    in_neighbors = node.get_in_neighbors()

    assert set(in_neighbors) == {1, 4}


# =============================================================================
# Category 4: DecentralizedCluster Topology Support (6 tests)
# =============================================================================

# 4.1 Cluster with Topology


@pytest.mark.asyncio
async def test_decentralizedcluster_nodes_with_topology():
    """Verify cluster nodes can be created with topology."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    topology = Topology.ring(3, k=1)
    cluster = DecentralizedCluster()

    for i in range(3):
        app = make_app(f"app-{i}")
        await cluster.add_node(node_id=i, application=app, topology=topology)

    # Verify nodes have correct topology
    for i in range(3):
        node = cluster.get_node(i)
        assert node.message_router.topology is topology

    await cluster.shutdown_all()


@pytest.mark.asyncio
async def test_decentralizedcluster_enforces_topology():
    """Verify cluster nodes respect topology constraints."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    topology = Topology.ring(5, k=1)
    cluster = DecentralizedCluster()

    for i in range(5):
        app = make_app(f"app-{i}")
        await cluster.add_node(node_id=i, application=app, topology=topology)

    await cluster.start_all()

    node0 = cluster.get_node(0)

    # Node 0 can send to 1 (neighbor)
    await node0.send_message(to_node_id=1, message_type="test", payload={})

    # Node 0 cannot send to 2 (not a neighbor)
    with pytest.raises(ValueError):
        await node0.send_message(to_node_id=2, message_type="test", payload={})

    await cluster.shutdown_all()


# 4.2 Cluster Broadcast via Node


@pytest.mark.asyncio
async def test_decentralizedcluster_node_sends_message():
    """Verify node in cluster can send messages to neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    # Use string node_ids that map to topology indices
    topology = Topology.complete(3)
    cluster = DecentralizedCluster()

    received = {"node0": [], "node1": [], "node2": []}

    for i in range(3):
        app = make_app(f"app-{i}")
        await cluster.add_node(node_id=f"node{i}", application=app, topology=topology)

    await cluster.start_all()

    # Register handlers AFTER start_all (like M3 tests)
    for node_id, node in cluster.nodes.items():
        async def handler(from_id, payload, nid=node_id):
            received[nid].append((from_id, payload))
        node.register_message_handler("test_msg", handler)

    # Send from node0 to node1 via the node's send_message
    node0 = cluster.get_node("node0")
    await node0.send_message(to_node_id="node1", message_type="test_msg", payload={"value": 42})

    await asyncio.sleep(0.5)

    assert len(received["node1"]) == 1
    assert received["node1"][0] == ("node0", {"value": 42})

    await cluster.shutdown_all()


@pytest.mark.asyncio
async def test_decentralizedcluster_broadcast_from_node():
    """Verify node can broadcast to all neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    # Use string node_ids for ProcessContext compatibility
    topology = Topology.ring(4, k=1)
    cluster = DecentralizedCluster()

    received_count = {"node0": 0, "node1": 0, "node2": 0, "node3": 0}

    for i in range(4):
        app = make_app(f"app-{i}")
        await cluster.add_node(node_id=f"node{i}", application=app, topology=topology)

    await cluster.start_all()

    # Register handlers AFTER start_all (like M3 tests)
    for node_id, node in cluster.nodes.items():
        async def handler(from_id, payload, nid=node_id):
            received_count[nid] += 1
        node.register_message_handler("broadcast", handler)

    # Node0 broadcasts (neighbors: node1, node3 in ring)
    node0 = cluster.get_node("node0")
    await node0.broadcast_message(message_type="broadcast", payload={})

    await asyncio.sleep(0.5)

    # Neighbors receive
    assert received_count["node1"] == 1
    assert received_count["node3"] == 1
    # Non-neighbor doesn't
    assert received_count["node2"] == 0

    await cluster.shutdown_all()


# 4.3 Process-Based Topology Routing


@pytest.mark.asyncio
async def test_decentralizedcluster_topology_with_processes():
    """Verify topology works with process-based nodes."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster
    from byzpy.engine.node.context import ProcessContext

    topology = Topology.ring(3, k=1)
    cluster = DecentralizedCluster()

    for i in range(3):
        app = make_app(f"app-{i}")
        await cluster.add_node(node_id=i, application=app, topology=topology)

    await cluster.start_all()

    # Verify processes are running with topology
    for i in range(3):
        node = cluster.get_node(i)
        assert isinstance(node.context, ProcessContext)
        assert node.context._process.is_alive()
        assert node.message_router.topology is topology

    await cluster.shutdown_all()


@pytest.mark.asyncio
async def test_decentralizedcluster_cross_process_topology_enforcement():
    """Verify topology enforcement works across processes with message handlers."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    # Use string node_ids for ProcessContext compatibility
    topology = Topology.ring(3, k=1)  # Node0 neighbors: node1, node2
    cluster = DecentralizedCluster()

    received = {"node0": [], "node1": [], "node2": []}

    for i in range(3):
        app = make_app(f"app-{i}")
        await cluster.add_node(node_id=f"node{i}", application=app, topology=topology)

    await cluster.start_all()

    # Register handlers AFTER start_all (like M3 tests)
    for node_id, node in cluster.nodes.items():
        async def handler(from_id, payload, nid=node_id):
            received[nid].append({"from": from_id, "data": payload})
        node.register_message_handler("gradient", handler)

    # Node0 sends to neighbor node1 (should succeed)
    node0 = cluster.get_node("node0")
    await node0.send_message(to_node_id="node1", message_type="gradient", payload={"v": 1})

    await asyncio.sleep(0.5)

    # Verify node1 received
    assert len(received["node1"]) == 1
    assert received["node1"][0]["from"] == "node0"

    await cluster.shutdown_all()


# =============================================================================
# Category 5: Topology Patterns (5 tests)
# =============================================================================

# 5.1 Ring Topology


@pytest.mark.asyncio
async def test_ring_topology_k1():
    """Verify ring topology with k=1 has correct neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.router import MessageRouter

    topology = Topology.ring(5, k=1)

    # Each node has 2 neighbors (prev and next)
    for i in range(5):
        router = MessageRouter(topology=topology, node_id=i)
        neighbors = router.get_out_neighbors()
        assert len(neighbors) == 2
        # Neighbors are (i-1) mod 5 and (i+1) mod 5
        assert (i - 1) % 5 in neighbors
        assert (i + 1) % 5 in neighbors


@pytest.mark.asyncio
async def test_ring_topology_k2():
    """Verify ring topology with k=2 has 4 neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.router import MessageRouter

    topology = Topology.ring(6, k=2)

    router = MessageRouter(topology=topology, node_id=0)
    neighbors = router.get_out_neighbors()

    # k=2: neighbors are ±1, ±2 from node 0
    assert set(neighbors) == {1, 2, 4, 5}


# 5.2 Complete Topology


@pytest.mark.asyncio
async def test_complete_topology():
    """Verify complete topology allows all-to-all communication."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.router import MessageRouter

    n = 5
    topology = Topology.complete(n)

    for i in range(n):
        router = MessageRouter(topology=topology, node_id=i)
        neighbors = router.get_out_neighbors()

        # Every node can send to all other nodes
        expected = set(range(n)) - {i}
        assert set(neighbors) == expected


# 5.3 Grid Topology (New)


@pytest.mark.asyncio
async def test_grid_topology():
    """Verify grid topology has correct neighbors."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.router import MessageRouter

    # Create 3x3 grid topology
    # Node layout:
    # 0 1 2
    # 3 4 5
    # 6 7 8

    def create_grid(rows, cols):
        n = rows * cols
        edges = []
        for r in range(rows):
            for c in range(cols):
                node = r * cols + c
                # Right neighbor
                if c < cols - 1:
                    edges.append((node, node + 1))
                    edges.append((node + 1, node))
                # Down neighbor
                if r < rows - 1:
                    edges.append((node, node + cols))
                    edges.append((node + cols, node))
        return Topology(n, edges)

    topology = create_grid(3, 3)

    # Node 0 (top-left) has neighbors: 1 (right), 3 (down)
    router0 = MessageRouter(topology=topology, node_id=0)
    assert set(router0.get_out_neighbors()) == {1, 3}

    # Node 4 (center) has neighbors: 1, 3, 5, 7
    router4 = MessageRouter(topology=topology, node_id=4)
    assert set(router4.get_out_neighbors()) == {1, 3, 5, 7}

    # Node 8 (bottom-right) has neighbors: 5, 7
    router8 = MessageRouter(topology=topology, node_id=8)
    assert set(router8.get_out_neighbors()) == {5, 7}


# =============================================================================
# Category 6: Error Handling (4 tests)
# =============================================================================


def test_messagerouter_invalid_node_id():
    """Verify MessageRouter rejects invalid node_id."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.router import MessageRouter

    topology = Topology.ring(3, k=1)

    # Node ID 10 doesn't exist in 3-node topology
    with pytest.raises(ValueError, match="not in topology"):
        MessageRouter(topology=topology, node_id=10)


@pytest.mark.asyncio
async def test_messagerouter_route_to_self_raises():
    """Verify routing to self raises error."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.router import MessageRouter

    topology = Topology.complete(3)
    router = MessageRouter(topology=topology, node_id=0)

    with pytest.raises(ValueError, match="cannot send to self"):
        await router.route_direct(
            target_node_id=0,
            message_type="test",
            payload={},
            context=None,
        )


def test_messagerouter_empty_topology():
    """Verify MessageRouter handles single-node topology."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.router import MessageRouter

    # Single node = no edges
    topology = Topology(1, [])
    router = MessageRouter(topology=topology, node_id=0)

    # No neighbors
    assert router.get_out_neighbors() == []
    assert router.get_in_neighbors() == []


@pytest.mark.asyncio
async def test_messagerouter_broadcast_no_neighbors():
    """Verify broadcast with no neighbors completes without error."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.router import MessageRouter

    topology = Topology(1, [])  # Single node
    router = MessageRouter(topology=topology, node_id=0)

    # Should not raise, just do nothing
    await router.route_broadcast(
        message_type="test",
        payload={},
        context=None,
    )


# =============================================================================
# Category 7: Integration with Existing Components (4 tests)
# =============================================================================


@pytest.mark.asyncio
async def test_topology_gradient_exchange_ring():
    """Verify gradient exchange follows ring topology constraints."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    # Use string node_ids for ProcessContext compatibility
    topology = Topology.ring(4, k=1)
    cluster = DecentralizedCluster()

    gradients_received = {"node0": [], "node1": [], "node2": [], "node3": []}

    # Register handlers BEFORE start_all (per test plan)
    for i in range(4):
        app = make_app(f"app-{i}")
        node = await cluster.add_node(node_id=f"node{i}", application=app, topology=topology)

        # Capture i in closure properly using a factory function
        def make_handler(nid):
            async def on_gradient(from_id, payload):
                gradients_received[nid].append(from_id)
            return on_gradient

        node.register_message_handler("gradient", make_handler(f"node{i}"))

    await cluster.start_all()

    # Each node broadcasts gradient to neighbors
    for node_id, node in cluster.nodes.items():
        await node.broadcast_message(
            message_type="gradient",
            payload={"grad": [1.0]},
        )

    await asyncio.sleep(0.5)

    # Verify topology is respected:
    # - Each node receives from exactly 2 neighbors (ring k=1)
    # - node0 neighbors: node1, node3
    # - node1 neighbors: node0, node2
    # - etc.
    for i in range(4):
        node_id = f"node{i}"
        assert len(gradients_received[node_id]) == 2
        expected = {f"node{(i - 1) % 4}", f"node{(i + 1) % 4}"}
        assert set(gradients_received[node_id]) == expected

    await cluster.shutdown_all()


@pytest.mark.asyncio
async def test_topology_aggregation_from_in_neighbors():
    """Verify node correctly identifies in-neighbors for aggregation."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.decentralized import DecentralizedNode
    from byzpy.engine.node.context import InProcessContext

    topology = Topology.ring(4, k=1)

    contexts = [InProcessContext() for _ in range(4)]
    apps = [make_app(f"app-{i}") for i in range(4)]

    nodes = [
        DecentralizedNode(node_id=i, application=apps[i], context=contexts[i], topology=topology)
        for i in range(4)
    ]

    for node in nodes:
        await node.start()

    # Node 0 should only receive from in-neighbors (1 and 3)
    expected_in_neighbors = nodes[0].get_in_neighbors()
    assert set(expected_in_neighbors) == {1, 3}

    # Node 2 should receive from (1 and 3)
    expected_in_neighbors_2 = nodes[2].get_in_neighbors()
    assert set(expected_in_neighbors_2) == {1, 3}

    for node in nodes:
        await node.shutdown()


@pytest.mark.asyncio
async def test_topology_ring_p2p_training_round():
    """Verify complete P2P training round with ring topology and aggregation."""
    import torch
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    # Use string node_ids for ProcessContext compatibility
    topology = Topology.ring(4, k=1)
    cluster = DecentralizedCluster()

    # Track state
    node_states = {f"node{i}": {"local_grad": None, "received": []} for i in range(4)}

    for i in range(4):
        app = make_app(f"app-{i}")
        await cluster.add_node(node_id=f"node{i}", application=app, topology=topology)

    await cluster.start_all()

    # Register handlers AFTER start_all
    for node_id, node in cluster.nodes.items():
        async def on_gradient(from_id, payload, nid=node_id):
            grad = torch.tensor(payload["grad"])
            node_states[nid]["received"].append(grad)
        node.register_message_handler("gradient", on_gradient)

    # Phase 1: Each node computes local gradient
    for node_id in node_states:
        node_states[node_id]["local_grad"] = torch.randn(5)

    # Phase 2: Each node broadcasts to neighbors
    for node_id, node in cluster.nodes.items():
        await node.broadcast_message(
            message_type="gradient",
            payload={"grad": node_states[node_id]["local_grad"].tolist()},
        )

    await asyncio.sleep(0.5)

    # Phase 3: Verify and aggregate
    for node_id in node_states:
        # Each node receives from 2 neighbors
        assert len(node_states[node_id]["received"]) == 2

        # Aggregate: mean of local + received
        all_grads = [node_states[node_id]["local_grad"]] + node_states[node_id]["received"]
        aggregated = torch.stack(all_grads).mean(dim=0)
        assert aggregated.shape == (5,)

    await cluster.shutdown_all()


@pytest.mark.asyncio
async def test_topology_complete_all_to_all():
    """Verify complete topology allows all-to-all communication."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    # Use string node_ids for ProcessContext compatibility
    topology = Topology.complete(3)
    cluster = DecentralizedCluster()

    received = {"node0": set(), "node1": set(), "node2": set()}

    for i in range(3):
        app = make_app(f"app-{i}")
        await cluster.add_node(node_id=f"node{i}", application=app, topology=topology)

    await cluster.start_all()

    # Register handlers AFTER start_all
    for node_id, node in cluster.nodes.items():
        async def on_msg(from_id, payload, nid=node_id):
            received[nid].add(from_id)
        node.register_message_handler("sync", on_msg)

    # Each node broadcasts to all others
    for node_id, node in cluster.nodes.items():
        await node.broadcast_message(message_type="sync", payload={})

    await asyncio.sleep(0.5)

    # Each node should receive from all other nodes
    for i in range(3):
        node_id = f"node{i}"
        expected = {f"node{j}" for j in range(3) if j != i}
        assert received[node_id] == expected

    await cluster.shutdown_all()


# =============================================================================
# Category 8: Multi-Node P2P Training Example (2 tests)
# =============================================================================


@pytest.mark.asyncio
async def test_ring_p2p_training_simulation():
    """End-to-end test: Ring topology P2P gradient exchange."""
    import torch
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    # Use string node_ids for ProcessContext compatibility
    n_nodes = 4
    topology = Topology.ring(n_nodes, k=1)

    cluster = DecentralizedCluster()

    # Track state per node
    node_states = {f"node{i}": {"gradients": [], "aggregated": None} for i in range(n_nodes)}

    for i in range(n_nodes):
        app = make_app(f"node-{i}")
        await cluster.add_node(node_id=f"node{i}", application=app, topology=topology)

    await cluster.start_all()

    # Register handlers AFTER start_all
    for node_id, node in cluster.nodes.items():
        # Capture node_id in closure properly using a factory function
        def make_handler(nid):
            async def on_gradient(from_id, payload):
                node_states[nid]["gradients"].append(payload["gradient"])
            return on_gradient

        node.register_message_handler("gradient", make_handler(node_id))

    # Phase 1: Each node computes and broadcasts local gradient
    for node_id, node in cluster.nodes.items():
        local_gradient = torch.randn(5)
        await node.broadcast_message(
            message_type="gradient",
            payload={"gradient": local_gradient.tolist()},
        )

    await asyncio.sleep(0.5)

    # Phase 2: Verify each node received from exactly 2 neighbors
    for node_id in node_states:
        assert len(node_states[node_id]["gradients"]) == 2

    await cluster.shutdown_all()


@pytest.mark.asyncio
@pytest.mark.skip(reason="Can run as standalone test, but not as part of the test suite")
async def test_complete_topology_consensus():
    """End-to-end test: Complete topology achieves consensus value."""
    from byzpy.engine.peer_to_peer.topology import Topology
    from byzpy.engine.node.cluster import DecentralizedCluster

    # Use string node_ids for ProcessContext compatibility
    n_nodes = 3
    topology = Topology.complete(n_nodes)

    cluster = DecentralizedCluster()

    values_received = {f"node{i}": [] for i in range(n_nodes)}
    initial_values = {"node0": 1.0, "node1": 2.0, "node2": 3.0}

    # Register handlers BEFORE start_all (per test plan)
    for i in range(n_nodes):
        app = make_app(f"node-{i}")
        node = await cluster.add_node(node_id=f"node{i}", application=app, topology=topology)

        # Capture i in closure properly using a factory function
        def make_handler(nid):
            async def on_value(from_id, payload):
                values_received[nid].append(payload["value"])
            return on_value

        node.register_message_handler("value", make_handler(f"node{i}"))

    await cluster.start_all()

    # Each node broadcasts its initial value
    for node_id, node in cluster.nodes.items():
        await node.broadcast_message(
            message_type="value",
            payload={"value": initial_values[node_id]},
        )

    await asyncio.sleep(0.3)

    # In complete topology, each node receives from all others
    for node_id in values_received:
        assert len(values_received[node_id]) == n_nodes - 1
        # Received values are from all other nodes
        expected_values = [initial_values[nid] for nid in initial_values if nid != node_id]
        assert sorted(values_received[node_id]) == sorted(expected_values)

    await cluster.shutdown_all()

