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

"""Unit tests for SequentialOffloadBackend."""

import pytest
import torch
from torch import nn

from vllm_omni.diffusion.hooks import HookRegistry
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
from vllm_omni.diffusion.offloader.sequential_backend import (
    ModelLevelOffloadBackend,
    SequentialOffloadHook,
    apply_sequential_offload,
    remove_sequential_offload,
    sequential_offload_component,
)
from vllm_omni.platforms import current_omni_platform

pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]


@pytest.fixture
def accelerator_device() -> torch.device:
    """Fixture that provides accelerator device or skips test if unavailable."""
    if current_omni_platform.get_device_count() == 0:
        pytest.skip("Accelerator required for this test")
    return current_omni_platform.get_torch_device(0)


def _create_simple_module() -> nn.Module:
    class SimpleModule(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(10, 20)

    return SimpleModule()


def _track_pin_memory_calls():
    tracker = {"called": False}
    original = torch.Tensor.pin_memory

    def mock(self):
        tracker["called"] = True
        return original(self)

    return tracker, mock


class _CustomPipeline(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.enable_args: dict[str, object] | None = None
        self.disable_called = False

    def enable_omni_model_cpu_offload(self, **kwargs) -> None:
        self.enable_args = kwargs

    def disable_omni_model_cpu_offload(self) -> None:
        self.disable_called = True


def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
    pipeline = _CustomPipeline()
    backend = ModelLevelOffloadBackend(
        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
        torch.device("cpu"),
    )

    backend.enable(pipeline)

    assert backend.enabled is True
    assert pipeline.enable_args == {
        "device": torch.device("cpu"),
        "pin_memory": False,
        "use_hsdp": False,
    }

    backend.disable()

    assert backend.enabled is False
    assert pipeline.disable_called is True


def test_model_level_backend_passes_explicit_component_selection() -> None:
    pipeline = _CustomPipeline()
    backend = ModelLevelOffloadBackend(
        OffloadConfig(
            strategy=OffloadStrategy.MODEL_LEVEL,
            components=frozenset({"dit", "text_encoder"}),
        ),
        torch.device("cpu"),
    )

    backend.enable(pipeline)

    assert pipeline.enable_args is not None
    assert pipeline.enable_args["offload_components"] == frozenset({"dit", "text_encoder"})


def test_model_level_backend_rolls_back_partial_custom_enable() -> None:
    class CustomPipeline(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.disable_called = False

        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
            del kwargs
            raise RuntimeError("injected custom enable failure")

        def disable_omni_model_cpu_offload(self) -> None:
            self.disable_called = True

    pipeline = CustomPipeline()
    backend = ModelLevelOffloadBackend(
        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL),
        torch.device("cpu"),
    )

    with pytest.raises(RuntimeError, match="injected custom enable failure"):
        backend.enable(pipeline)

    assert pipeline.disable_called
    assert not backend.enabled


def test_sequential_offload_rolls_back_partial_hook_registration(monkeypatch: pytest.MonkeyPatch) -> None:
    dit = _create_simple_module()
    encoder = _create_simple_module()
    original_register = HookRegistry.register_hook
    calls = 0

    def fail_second_registration(self, name, hook):
        nonlocal calls
        calls += 1
        if calls == 2:
            raise RuntimeError("injected registration failure")
        return original_register(self, name, hook)

    monkeypatch.setattr(HookRegistry, "register_hook", fail_second_registration)

    with pytest.raises(RuntimeError, match="injected registration failure"):
        apply_sequential_offload(
            dit_modules=[dit],
            encoder_modules=[encoder],
            device=torch.device("cpu"),
        )

    for module in (dit, encoder):
        registry = getattr(module, "_hook_registry", None)
        assert registry is None or registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None


def test_model_level_disable_cleans_remaining_hooks_and_can_retry(monkeypatch: pytest.MonkeyPatch) -> None:
    pipeline = nn.Module()
    pipeline.transformer = _create_simple_module()
    pipeline.text_encoder = _create_simple_module()
    backend = ModelLevelOffloadBackend(
        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
        torch.device("cpu"),
    )
    backend.enable(pipeline)

    original_remove = HookRegistry.remove_hook
    calls = 0

    def fail_first_removal(self, name):
        nonlocal calls
        calls += 1
        if calls == 1:
            raise RuntimeError("injected hook removal failure")
        original_remove(self, name)

    monkeypatch.setattr(HookRegistry, "remove_hook", fail_first_removal)

    with pytest.raises(RuntimeError, match="Failed to remove one or more sequential offload hooks"):
        backend.disable()

    assert backend.enabled
    assert backend._offload_modules == [pipeline.transformer, pipeline.text_encoder]
    assert pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is not None
    assert pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None

    backend.disable()

    assert not backend.enabled
    assert not backend._offload_modules
    for module in (pipeline.transformer, pipeline.text_encoder):
        assert module._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None


def test_sequential_offload_filters_cpu_eligible_components() -> None:
    dit = _create_simple_module()
    encoder = _create_simple_module()
    resident_stage = _create_simple_module()

    apply_sequential_offload(
        dit_modules=[dit],
        encoder_modules=[encoder, resident_stage],
        device=torch.device("cpu"),
        offload_dit_modules=[dit],
        offload_encoder_modules=[encoder],
    )

    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    resident_hook = resident_stage._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    assert dit_hook.offload_targets == [encoder]
    assert encoder_hook.offload_targets == [dit]
    assert encoder_hook.offload_after_context is True
    assert resident_hook.offload_targets == [dit]
    assert resident_hook.offload_after_context is False

    remove_sequential_offload([dit, encoder, resident_stage])


def test_move_params_handles_buffer_only_modules() -> None:
    module = nn.Module()
    module.register_buffer("state", torch.ones(2))

    # An indexed CPU device compares differently while remaining usable on
    # CPU-only test hosts, so the helper must inspect the buffer itself.
    moved = SequentialOffloadHook._move_params(module, torch.device("cpu:1"))

    assert moved
    assert module.state.device.type == "cpu"


def test_model_level_backend_keeps_declared_image_encoder_resident() -> None:
    class MixedEncoderPipeline(nn.Module):
        _offload_plan = OffloadPlan(
            encoder_component_types={
                "text_encoder": "text_encoder",
            }
        )

        def __init__(self) -> None:
            super().__init__()
            self.transformer = _create_simple_module()
            self.text_encoder = _create_simple_module()
            self.image_encoder = _create_simple_module()

    pipeline = MixedEncoderPipeline()
    backend = ModelLevelOffloadBackend(
        OffloadConfig(
            strategy=OffloadStrategy.MODEL_LEVEL,
            components=frozenset({"dit", "text_encoder"}),
        ),
        torch.device("cpu"),
    )

    backend.enable(pipeline)

    text_hook = pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    image_hook = pipeline.image_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    assert text_hook.offload_after_context is True
    assert image_hook.offload_after_context is False

    backend.disable()


@pytest.mark.parametrize(
    ("component", "attribute", "message"),
    [
        ("text_encoder", "text_encoder", "requires a DiT/transformer"),
        ("dit", "transformer", "requires an encoder execution stage"),
    ],
)
def test_component_selective_model_offload_requires_swap_counterpart(component, attribute, message) -> None:
    pipeline = nn.Module()
    setattr(pipeline, attribute, _create_simple_module())
    backend = ModelLevelOffloadBackend(
        OffloadConfig(
            strategy=OffloadStrategy.MODEL_LEVEL,
            components=frozenset({component}),
        ),
        torch.device("cpu"),
    )

    with pytest.raises(ValueError, match=message):
        backend.enable(pipeline)


def test_persistent_dit_staging_reuses_fixed_device_storage(accelerator_device) -> None:
    dit = _create_simple_module().to(accelerator_device)
    encoder = _create_simple_module().to(accelerator_device)
    original_weight = dit.linear.weight.detach().clone()

    apply_sequential_offload(
        dit_modules=[dit],
        encoder_modules=[encoder],
        device=accelerator_device,
        pin_memory=False,
        persistent_dit_staging=True,
    )

    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    assert dit_hook._stager is not None
    assert encoder_hook._stager is None  # encoders keep plain move semantics

    # Registration snapshots the CPU master and rebinds to it.
    assert dit.linear.weight.device.type == "cpu"

    dit_hook._to_gpu(dit)
    first_ptr = dit.linear.weight.data_ptr()
    assert dit.linear.weight.device == accelerator_device
    torch.testing.assert_close(dit.linear.weight.data, original_weight)

    # Offload via a *different* module's hook (the cross-module swap path),
    # then reload: the same fixed device storage must be reused so that
    # CUDA-graph-captured weight pointers stay valid.
    encoder_hook._to_cpu(dit)
    assert dit.linear.weight.device.type == "cpu"
    dit_hook._to_gpu(dit)
    assert dit.linear.weight.data_ptr() == first_ptr
    torch.testing.assert_close(dit.linear.weight.data, original_weight)

    # A second offload/load cycle through the DiT hook itself stays put too.
    dit_hook._to_cpu(dit)
    dit_hook._to_gpu(dit)
    assert dit.linear.weight.data_ptr() == first_ptr

    remove_sequential_offload([dit, encoder])


def test_default_sequential_offload_releases_device_storage(accelerator_device) -> None:
    """Without persistent staging, offloading keeps plain move semantics."""
    dit = _create_simple_module().to(accelerator_device)
    encoder = _create_simple_module().to(accelerator_device)

    apply_sequential_offload(
        dit_modules=[dit],
        encoder_modules=[encoder],
        device=accelerator_device,
        pin_memory=False,
    )

    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    assert dit_hook._stager is None  # no retained staging: GPU storage is freed on offload

    dit_hook._to_gpu(dit)
    assert dit.linear.weight.device == accelerator_device
    dit_hook._to_cpu(dit)
    assert dit.linear.weight.device.type == "cpu"

    remove_sequential_offload([dit, encoder])


@pytest.mark.parametrize(
    ("decode_graphs", "expect_stager"),
    [(True, True), (False, False)],
)
def test_model_level_persistent_staging_gated_on_decode_graph(
    accelerator_device, decode_graphs: bool, expect_stager: bool
) -> None:
    """ModelLevelOffloadBackend only retains DiT staging for decode-graph models."""
    pipeline = nn.Module()
    pipeline.transformer = _create_simple_module()
    if decode_graphs:
        pipeline.transformer.enable_cuda_graph_decode = True
    pipeline.text_encoder = _create_simple_module()
    backend = ModelLevelOffloadBackend(
        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
        accelerator_device,
    )
    backend.enable(pipeline)

    hook = pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    assert (hook._stager is not None) is expect_stager

    backend.disable()


def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
    dit = _create_simple_module()
    encoder = _create_simple_module()
    offloaded: list[nn.Module] = []
    monkeypatch.setattr(
        SequentialOffloadHook,
        "_to_cpu",
        lambda self, module: offloaded.append(module),
    )

    apply_sequential_offload(
        dit_modules=[dit],
        encoder_modules=[encoder],
        device=torch.device("cpu"),
        offload_initial_dits=True,
    )

    assert offloaded == [dit]
    remove_sequential_offload([dit, encoder])


def test_direct_component_activation_failure_still_offloads(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    component = _create_simple_module()
    apply_sequential_offload(
        dit_modules=[_create_simple_module()],
        encoder_modules=[component],
        device=torch.device("cpu"),
    )
    hook = component._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
    offloaded: list[nn.Module] = []

    def fail_activation(module: nn.Module) -> None:
        raise RuntimeError("activation failed")

    monkeypatch.setattr(hook, "pre_forward", fail_activation)
    monkeypatch.setattr(hook, "_to_cpu", lambda module: offloaded.append(module))

    with pytest.raises(RuntimeError, match="activation failed"):
        with sequential_offload_component(component):
            pass

    assert offloaded == [component]


class TestMoveParamsPinMemory:
    def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
        """DTensor should skip pin_memory to avoid RuntimeError."""
        module = _create_simple_module().to(accelerator_device)
        tracker, mock_pin = _track_pin_memory_calls()

        original_isinstance = isinstance

        def fake_isinstance(obj, cls):
            # torchada's Tensor.to() checks tuples such as
            # ``(str, torch.device)``.  Keep the DTensor probe compatible with
            # the normal isinstance contract instead of breaking those calls.
            if type(cls) is tuple:
                return any(fake_isinstance(obj, candidate) for candidate in cls)
            if getattr(cls, "__name__", None) == "DTensor":
                return True
            return original_isinstance(obj, cls)

        monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
        monkeypatch.setattr("builtins.isinstance", fake_isinstance)
        hook = SequentialOffloadHook(
            offload_targets=[],
            device=accelerator_device,
            pin_memory=True,
            use_hsdp=False,
        )
        hook._move_params(
            module,
            torch.device("cpu"),
            non_blocking=False,
            pin_memory=True,
        )
        assert not tracker["called"], "pin_memory should not be called for DTensor"

    def test_regular_tensor_calls_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
        """Regular tensor should call pin_memory when moving to CPU."""
        module = _create_simple_module().to(accelerator_device)
        tracker, mock_pin = _track_pin_memory_calls()

        monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
        hook = SequentialOffloadHook(
            offload_targets=[],
            device=accelerator_device,
            pin_memory=True,
            use_hsdp=False,
        )
        hook._move_params(
            module,
            torch.device("cpu"),
            non_blocking=False,
            pin_memory=True,
        )
        assert tracker["called"], "pin_memory should be called for regular tensors"

    def test_pin_memory_skipped_when_disabled(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
        """pin_memory should not be called when pin_memory=False."""
        module = _create_simple_module().to(accelerator_device)
        tracker, mock_pin = _track_pin_memory_calls()

        monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
        hook = SequentialOffloadHook(
            offload_targets=[],
            device=accelerator_device,
            pin_memory=False,
            use_hsdp=False,
        )
        hook._move_params(
            module,
            torch.device("cpu"),
            non_blocking=False,
            pin_memory=False,
        )
        assert not tracker["called"], "pin_memory should not be called when disabled"

    def test_pin_memory_skipped_for_non_cpu_target(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
        """pin_memory should not be called for non-CPU targets."""
        module = _create_simple_module().to("cpu")
        tracker, mock_pin = _track_pin_memory_calls()

        monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
        hook = SequentialOffloadHook(
            offload_targets=[],
            device=torch.device("cpu"),
            pin_memory=True,
            use_hsdp=False,
        )
        hook._move_params(module, accelerator_device, non_blocking=False, pin_memory=True)
        assert not tracker["called"], "pin_memory should not be called for non-CPU target"