# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import pytest
import torch
import torch.nn as nn

from vllm.model_executor.offloader.prefetch import (
    _get_next_prefetch_index,
    _TowerUnitStreamer,
)

pytestmark = pytest.mark.skip_global_cleanup


@pytest.mark.parametrize(
    ("module_count", "prefetch_step", "expected_targets"),
    [
        (3, 1, [1, 2, 0]),
        (3, 2, [2, 1, 0]),
        (4, 2, [2, 3, 0, 1]),
        (5, 2, [2, 3, 4, 1, 0]),
        (3, 3, [0, 1, 2]),
        (3, 4, [1, 2, 0]),
    ],
)
def test_next_prefetch_index_preserves_slot_ownership(
    module_count: int,
    prefetch_step: int,
    expected_targets: list[int],
) -> None:
    targets = [
        _get_next_prefetch_index(index, prefetch_step, module_count)
        for index in range(module_count)
    ]

    assert targets == expected_targets
    if prefetch_step < module_count:
        assert all(
            target % prefetch_step == index % prefetch_step
            for index, target in enumerate(targets)
        )


def test_tower_unit_streamer_roundtrip_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
    """Host-resident weights must compute identically to resident weights,
    and be back on host storage after each forward."""
    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)

    lin = nn.Linear(4, 4, bias=False)
    weight = lin.weight.detach().clone()

    streamer = _TowerUnitStreamer(lin, ["weight"], torch.device("cpu"))
    streamer.offload_to_host_now()
    assert lin.weight.device.type == "cpu"

    streamer.sync_storage()
    assert streamer.offloaded_bytes == weight.numel() * weight.element_size()

    streamer.hook_forward()
    x = torch.randn(2, 4)
    out = lin(x)
    assert torch.allclose(out, x @ weight.T)
    # Weights released back to host after the forward.
    assert lin.weight.device.type == "cpu"

    # Re-running post_init-style hooking must not double-wrap the forward.
    streamer.hook_forward()
    assert torch.allclose(lin(x), out)


def test_offload_config_tower_only_prefetch_is_valid() -> None:
    from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig

    cfg = OffloadConfig(
        offload_backend="prefetch",
        prefetch=PrefetchOffloadConfig(offload_towers=True),
    )
    assert cfg.prefetch.offload_towers
    assert cfg.prefetch.offload_group_size == 0


def test_offload_config_prefetch_requires_layers_or_towers() -> None:
    from vllm.config.offload import OffloadConfig

    with pytest.raises(ValueError):
        OffloadConfig(offload_backend="prefetch")


@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
def test_create_offloader_selects_prefetch_for_towers() -> None:
    from vllm.config.offload import OffloadConfig, PrefetchOffloadConfig
    from vllm.model_executor.offloader.base import create_offloader
    from vllm.model_executor.offloader.prefetch import PrefetchOffloader

    offloader = create_offloader(
        OffloadConfig(prefetch=PrefetchOffloadConfig(offload_towers=True))
    )
    assert isinstance(offloader, PrefetchOffloader)
    assert offloader.supports_tower_offload