import pytest
from types import SimpleNamespace

import vllm_ascend.worker.model_runner_v1 as model_runner_v1
from vllm_ascend.model_executor.offloader import AscendPrefetchOffloader
from vllm_ascend.worker.model_runner_v1 import create_ascend_weight_offloader


def make_config(group_size, num_in_group, offload_backend="auto"):
    return SimpleNamespace(
        offload_backend=offload_backend,
        prefetch=SimpleNamespace(
            offload_group_size=group_size,
            offload_num_in_group=num_in_group,
            offload_params=set(),
        )
    )


def test_uva_backend_preserves_parent_created_offloader(monkeypatch):
    parent_offloader = object()
    monkeypatch.setattr(model_runner_v1, "get_offloader", lambda: parent_offloader)

    offloader = create_ascend_weight_offloader(
        make_config(4, 1, offload_backend="uva")
    )

    assert offloader is parent_offloader


def test_auto_backend_with_zero_group_size_preserves_parent_created_offloader(
    monkeypatch,
):
    parent_offloader = object()
    monkeypatch.setattr(model_runner_v1, "get_offloader", lambda: parent_offloader)

    offloader = create_ascend_weight_offloader(
        make_config(0, 1, offload_backend="auto")
    )

    assert offloader is parent_offloader


def test_auto_backend_with_positive_group_size_replaces_parent_with_prefetch(
    monkeypatch,
):
    parent_offloader = object()
    monkeypatch.setattr(model_runner_v1, "get_offloader", lambda: parent_offloader)

    offloader = create_ascend_weight_offloader(
        make_config(4, 1, offload_backend="auto")
    )

    assert isinstance(offloader, AscendPrefetchOffloader)
    assert offloader is not parent_offloader
    assert offloader.group_size == 4
    assert offloader.num_in_group == 1
    assert offloader.prefetch_step == 1


def test_prefetch_backend_replaces_parent_created_offloader(monkeypatch):
    parent_offloader = object()
    monkeypatch.setattr(model_runner_v1, "get_offloader", lambda: parent_offloader)

    offloader = create_ascend_weight_offloader(
        make_config(4, 1, offload_backend="prefetch")
    )

    assert isinstance(offloader, AscendPrefetchOffloader)
    assert offloader is not parent_offloader


def test_negative_group_size_raises():
    with pytest.raises(ValueError, match="offload_group_size must be >= 0"):
        create_ascend_weight_offloader(make_config(-1, 1))


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