import pytest

from vllm.model_executor.offloader.base import NoopOffloader
from vllm_ascend.model_executor.offloader import AscendPrefetchOffloader
from vllm_ascend.model_executor.offloader.selection import select_offload_layers
from vllm_ascend.worker.model_runner_v1 import create_ascend_weight_offloader

from tests.model_executor.offloader.test_runner_integration import make_config


@pytest.mark.parametrize(
    ("num_in_group", "expected_layers"),
    [
        (0, ()),
        (1, (3, 7, 11)),
        (2, (2, 3, 6, 7, 10, 11)),
        (4, tuple(range(12))),
    ],
)
def test_offload_num_in_group_controls_layers_selected_within_each_group(
    num_in_group,
    expected_layers,
):
    selection = select_offload_layers(
        num_layers=12,
        group_size=4,
        num_in_group=num_in_group,
    )

    assert selection.layer_indices == expected_layers
    assert selection.enabled is bool(expected_layers)


def test_offload_num_in_group_zero_disables_offloader_even_when_group_size_set():
    offloader = create_ascend_weight_offloader(make_config(4, 0))

    assert isinstance(offloader, NoopOffloader)


def test_offload_num_in_group_larger_than_group_size_is_rejected():
    with pytest.raises(ValueError, match="offload_num_in_group"):
        create_ascend_weight_offloader(make_config(2, 3))


def test_valid_offload_num_in_group_is_preserved_on_prefetch_offloader():
    offloader = create_ascend_weight_offloader(make_config(4, 2))

    assert isinstance(offloader, AscendPrefetchOffloader)
    assert offloader.group_size == 4
    assert offloader.num_in_group == 2
