from __future__ import annotations

from rolemesh_gateway.registry import NodeHeartbeat, NodeRegistration, Registry


def test_registry_persists_round_robin_and_heartbeat_state(tmp_path):
    persist_path = tmp_path / "registry.json"
    registry = Registry(persist_path=persist_path)

    registry.register(
        NodeRegistration(
            node_id="node-a",
            base_url="http://127.0.0.1:9001",
            roles=["reviewer"],
        )
    )
    registry.register(
        NodeRegistration(
            node_id="node-b",
            base_url="http://127.0.0.1:9002",
            roles=["reviewer"],
        )
    )

    assert registry.pick_node_for_role("reviewer").node_id == "node-a"
    assert registry.pick_node_for_role("reviewer").node_id == "node-b"

    node = registry.heartbeat(
        NodeHeartbeat(
            node_id="node-a",
            timestamp=456.0,
            status={"healthy": True},
            metrics=[{"queue_depth": 1}],
        )
    )
    assert node is not None
    assert node.status["metrics"] == [{"queue_depth": 1}]

    reloaded = Registry(persist_path=persist_path)
    assert reloaded.pick_node_for_role("reviewer").node_id == "node-a"
    assert reloaded.list_nodes()[0].status["timestamp"] == 456.0


def test_registry_heartbeat_returns_none_for_unknown_node():
    registry = Registry()
    assert registry.heartbeat(NodeHeartbeat(node_id="missing")) is None