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

import pytest
import torch.nn as nn

import vllm_omni.diffusion.compile as compile_module
from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile

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


class _WrappedBlock(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.compile_called = False
        self.forward_compiled = False

    def compile(self, *args, **kwargs):
        self.compile_called = True
        return self

    def forward(self, x):
        return x


class _ModelWithWrappedRepeatedBlocks(nn.Module):
    _repeated_blocks = ["OriginalBlock"]
    _layerwise_offload_blocks_attrs = ["transformer_blocks"]

    def __init__(self) -> None:
        super().__init__()
        self.transformer_blocks = nn.ModuleList([_WrappedBlock(), _WrappedBlock()])
        self.other_blocks = nn.ModuleList([_WrappedBlock()])


def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(monkeypatch):
    model = _ModelWithWrappedRepeatedBlocks()
    compile_calls = []

    def _compile(fn, *args, **kwargs):
        compile_calls.append((fn, args, kwargs))

        def _compiled(*fn_args, **fn_kwargs):
            return f"compiled:{fn(*fn_args, **fn_kwargs)}"

        return _compiled

    monkeypatch.setattr(compile_module.torch, "compile", _compile)

    regionally_compile(model, dynamic=True)

    assert len(compile_calls) == 2
    assert all(not block.compile_called for block in model.transformer_blocks)
    assert not model.other_blocks[0].compile_called
    assert model.transformer_blocks[0].forward("ok") == "compiled:ok"


def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypatch):
    model = _ModelWithWrappedRepeatedBlocks()
    original_forwards = [block.forward.__func__ for block in model.transformer_blocks]
    compile_calls = 0

    def _compile(fn, *args, **kwargs):
        nonlocal compile_calls
        compile_calls += 1
        if compile_calls == 2:
            raise RuntimeError("compile setup failed")
        return lambda *fn_args, **fn_kwargs: fn(*fn_args, **fn_kwargs)

    monkeypatch.setattr(compile_module.torch, "compile", _compile)

    with pytest.raises(RuntimeError, match="compile setup failed"):
        regionally_compile(model, dynamic=True)

    assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards


def test_compile_callable_falls_back_once_after_lazy_failure(monkeypatch, caplog):
    compiled_calls = 0
    eager_calls = 0

    def eager(value):
        nonlocal eager_calls
        eager_calls += 1
        return value + 1

    def compile_fn(fn, **kwargs):
        assert fn is eager
        assert kwargs == {"mode": "default"}

        def compiled(value):
            nonlocal compiled_calls
            compiled_calls += 1
            raise RuntimeError("lazy failure")

        return compiled

    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable", mode="default")

    assert activated
    assert callable_fn(1) == 2
    assert callable_fn(2) == 3
    assert compiled_calls == 1
    assert eager_calls == 2
    assert "disabling compilation and retrying eagerly" in caplog.text


def test_compile_callable_keeps_eager_on_setup_failure(monkeypatch, caplog):
    def eager(value):
        return value + 1

    def compile_fn(*args, **kwargs):
        raise RuntimeError("setup failure")

    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable")

    assert not activated
    assert callable_fn is eager
    assert callable_fn(1) == 2
    assert "continuing eagerly" in caplog.text