import torch
import torch.nn as nn

from vllm_ascend.model_executor.offloader.sync import AscendSyncOffloader


class TinyLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 2, bias=False)

    def forward(self, x):
        return self.linear(x)


def test_disabled_config_returns_modules_unchanged():
    modules = [TinyLayer(), TinyLayer()]
    offloader = AscendSyncOffloader(group_size=0, num_in_group=1)

    wrapped = offloader.wrap_modules(iter(modules))

    assert wrapped == modules
    assert offloader.selected_layer_indices == ()


def test_selects_matching_layers_without_touching_unselected_layers():
    modules = [TinyLayer() for _ in range(4)]
    offloader = AscendSyncOffloader(group_size=4, num_in_group=1)

    wrapped = offloader.wrap_modules(iter(modules))

    assert wrapped == modules
    assert offloader.selected_layer_indices == (3,)
    assert len(offloader.module_offloaders) == 1
    assert offloader.module_offloaders[0].layer_idx == 3


def test_invalid_num_in_group_raises():
    offloader = AscendSyncOffloader(group_size=2, num_in_group=3)

    try:
        offloader.wrap_modules(iter([TinyLayer(), TinyLayer()]))
    except ValueError as exc:
        assert "offload_num_in_group" in str(exc)
    else:
        raise AssertionError("expected ValueError")


def test_post_init_syncs_loaded_weights_to_cpu_storage():
    layer = torch.nn.Linear(2, 2, bias=False)
    modules = [torch.nn.Linear(2, 2, bias=False) for _ in range(3)] + [layer]
    offloader = AscendSyncOffloader(group_size=4, num_in_group=1)

    offloader.wrap_modules(iter(modules))
    loaded_weight = torch.full_like(layer.weight.data, 3.0)
    layer.weight.data = loaded_weight

    offloader.post_init()

    assert len(offloader.module_offloaders) == 1
    state = offloader.module_offloaders[0].param_states[0]
    assert state.param.device.type == "cpu"
    assert torch.equal(state.cpu_tensor, loaded_weight.cpu())
    assert offloader.total_offloaded_bytes == loaded_weight.numel() * loaded_weight.element_size()
