"""Tests for adaptive collective selector (Task1.6)."""

from __future__ import annotations

from sagellm_comm.benchmark.adaptive_selector import (
    AdaptiveCollectiveSelector,
    SelectorPolicyConfig,
)


class TestAdaptiveCollectiveSelector:
    """Validate backend/algo selection and traceability fields."""

    def test_auto_select_cpu_prefers_gloo_and_ring(self):
        selector = AdaptiveCollectiveSelector()

        decision = selector.select(
            world_size=2,
            topology=None,
            message_bytes=1024,
            requested_backend="auto",
            requested_algo="auto",
            available_backends=["gloo", "nccl"],
        )

        assert decision.backend_kind == "gloo"
        assert decision.algo_id == "ring"
        assert "auto_backend" in decision.rule_id
        assert "auto_algo" in decision.rule_id
        assert decision.reason

    def test_auto_select_tree_for_large_world(self):
        selector = AdaptiveCollectiveSelector()

        decision = selector.select(
            world_size=16,
            topology=None,
            message_bytes=1024,
            requested_backend="auto",
            requested_algo="auto",
            available_backends=["gloo"],
        )

        assert decision.algo_id == "tree"
        assert "auto_algo_tree_large_world" in decision.rule_id

    def test_auto_select_hybrid_for_large_message(self):
        selector = AdaptiveCollectiveSelector(
            SelectorPolicyConfig(
                backend_priority_by_device={"cpu": ["gloo"], "default": ["gloo"]},
                tree_world_size_threshold=64,
                hybrid_bytes_threshold=4096,
            )
        )

        decision = selector.select(
            world_size=8,
            topology=None,
            message_bytes=8192,
            requested_backend="auto",
            requested_algo="auto",
            available_backends=["gloo"],
        )

        assert decision.algo_id == "hybrid"
        assert "auto_algo_hybrid_large_msg" in decision.rule_id

    def test_explicit_backend_and_algo(self):
        selector = AdaptiveCollectiveSelector()

        decision = selector.select(
            world_size=4,
            topology=None,
            message_bytes=2048,
            requested_backend="nccl",
            requested_algo="tree",
            available_backends=["gloo", "nccl"],
        )

        assert decision.backend_kind == "nccl"
        assert decision.algo_id == "tree"
        assert "explicit_backend" in decision.rule_id
        assert "explicit_algo" in decision.rule_id
