from __future__ import annotations

import random

from rolemesh_gateway.registry import NodeHeartbeat, NodeRegistration, Registry, ServedModel


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"


def test_registry_supports_served_models_and_candidate_selection():
    registry = Registry()
    registry.register(
        NodeRegistration(
            node_id="node-a",
            base_url="http://127.0.0.1:9001",
            served_models=[
                ServedModel(model_id="qwen3-8b", roles=["tutor", "mentor"]),
                ServedModel(model_id="qwen2.5-coder-14b", roles=["code_tutor"]),
            ],
        )
    )

    tutor_candidates = registry.candidates_for_role("tutor")
    assert len(tutor_candidates) == 1
    assert tutor_candidates[0].node_id == "node-a"
    assert tutor_candidates[0].model_id == "qwen3-8b"

    mentor_candidates = registry.candidates_for_role("mentor")
    assert len(mentor_candidates) == 1
    assert mentor_candidates[0].model_id == "qwen3-8b"

    code_candidates = registry.candidates_for_role("code_tutor")
    assert len(code_candidates) == 1
    assert code_candidates[0].model_id == "qwen2.5-coder-14b"

    picked = registry.pick_candidate_for_role("mentor")
    assert picked is not None
    assert picked.model_id == "qwen3-8b"


def test_registry_loads_legacy_roles_as_served_models(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"],
        )
    )

    reloaded = Registry(persist_path=persist_path)
    candidates = reloaded.candidates_for_role("reviewer")
    assert len(candidates) == 1
    assert candidates[0].model_id == "reviewer"