from __future__ import annotations

import random

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


def test_registry_filters_stale_nodes_from_routing(tmp_path):
    persist_path = tmp_path / "registry.json"
    registry = Registry(persist_path=persist_path, stale_after_s=0.01)

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

    stale.last_seen = stale.last_seen - 60
    registry._save()

    assert registry.is_stale(stale) is True
    assert registry.is_stale(fresh) is False
    assert [node.node_id for node in registry.nodes_for_role("reviewer", include_stale=False)] == ["fresh-node"]
    assert registry.pick_node_for_role("reviewer").node_id == "fresh-node"


def test_registry_supports_random_selection(monkeypatch):
    registry = Registry()
    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"],
        )
    )

    monkeypatch.setattr(random, "choice", lambda candidates: candidates[-1])

    picked = registry.pick_node_for_role("reviewer", strategy="random")
    assert picked is not None
    assert picked.node_id == "node-b"