import torch

import vllm_ascend.model_executor.offloader.prefetch as prefetch_module
from vllm_ascend.model_executor.offloader import AscendPrefetchOffloader


def make_layers(n):
    return [torch.nn.Linear(2, 2, bias=False) for _ in range(n)]


class StridedWeightLayer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.empty_strided((2, 3), (1, 2)))


def make_layers_with_strided_selected_weight():
    return [torch.nn.Linear(2, 2, bias=False) for _ in range(3)] + [
        StridedWeightLayer()
    ]


def test_prefetch_wraps_selected_layers_and_records_slots():
    layers = make_layers(8)
    offloader = AscendPrefetchOffloader(
        group_size=4,
        num_in_group=1,
        prefetch_step=1,
    )

    wrapped = offloader.wrap_modules(iter(layers))

    assert wrapped == layers
    assert offloader.selected_layer_indices == (3, 7)
    assert [m.layer_idx for m in offloader.module_offloaders] == [3, 7]
    assert [m.slot_idx for m in offloader.module_offloaders] == [0, 0]


def test_prefetch_captures_pinned_cpu_storage_with_original_layout(monkeypatch):
    monkeypatch.setattr(
        prefetch_module,
        "is_pin_memory_available",
        lambda: True,
        raising=False,
    )
    layers = make_layers_with_strided_selected_weight()
    selected = layers[3]
    selected.weight.data.fill_(2.0)
    expected_shape = selected.weight.shape
    expected_stride = selected.weight.stride()
    expected_dtype = selected.weight.dtype
    offloader = AscendPrefetchOffloader(
        group_size=4,
        num_in_group=1,
        prefetch_step=1,
    )

    offloader.wrap_modules(iter(layers))

    state = offloader.module_offloaders[0].param_states[0]
    assert state.cpu_tensor.is_pinned()
    assert state.cpu_tensor.device.type == "cpu"
    assert state.cpu_tensor.shape == expected_shape
    assert state.cpu_tensor.stride() == expected_stride
    assert state.cpu_tensor.dtype == expected_dtype
    assert torch.equal(state.cpu_tensor, torch.full_like(state.cpu_tensor, 2.0))


def test_prefetch_post_init_keeps_cpu_storage_and_counts_bytes_without_npu():
    layers = make_layers(4)
    selected = layers[3]
    selected.weight.data.fill_(2.0)
    offloader = AscendPrefetchOffloader(
        group_size=4,
        num_in_group=1,
        prefetch_step=1,
    )
    offloader.wrap_modules(iter(layers))

    offloader.post_init()

    module_offloader = offloader.module_offloaders[0]
    state = module_offloader.param_states[0]
    assert state.cpu_tensor.device.type == "cpu"
    assert torch.equal(state.cpu_tensor, torch.full_like(state.cpu_tensor, 2.0))
    assert offloader.total_offloaded_bytes == selected.weight.numel() * selected.weight.element_size()


def test_prefetch_post_init_repins_processed_cpu_weights_when_available(monkeypatch):
    monkeypatch.setattr(
        prefetch_module,
        "is_pin_memory_available",
        lambda: True,
        raising=False,
    )
    layers = make_layers(4)
    selected = layers[3]
    offloader = AscendPrefetchOffloader(
        group_size=4,
        num_in_group=1,
        prefetch_step=1,
    )
    offloader.wrap_modules(iter(layers))

    loaded_weight = torch.full_like(selected.weight.data, 3.0)
    selected.weight.data = loaded_weight
    offloader.post_init()

    state = offloader.module_offloaders[0].param_states[0]
    assert state.cpu_tensor.is_pinned()
    assert state.cpu_tensor.device.type == "cpu"
    assert state.cpu_tensor.shape == loaded_weight.shape
    assert state.cpu_tensor.stride() == loaded_weight.stride()
    assert state.cpu_tensor.dtype == loaded_weight.dtype
    assert torch.equal(state.cpu_tensor, loaded_weight)


def test_prefetch_zero_step_is_rejected():
    try:
        AscendPrefetchOffloader(group_size=4, num_in_group=1, prefetch_step=0)
    except ValueError as exc:
        assert "prefetch_step" in str(exc)
    else:
        raise AssertionError("expected ValueError")
