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

from __future__ import annotations

import pytest
import torch

from tests.diffusion.models.magi2.test_native_packing import (
    _longer_text_tensors,
    _sampler_tensors,
    _tiny_model,
    _tiny_sampler,
)
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig

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


def test_layer_regions_bracket_the_eager_kernels() -> None:
    moe_layer, dense_layer = _tiny_model().block.layers

    assert moe_layer.region_methods == ("_attention_input", "_moe_input", "_moe_output")
    assert dense_layer.region_methods == ("_attention_input", "_dense_output")


def test_compile_regions_fullgraph_matches_eager_bitwise() -> None:
    model = _tiny_model()
    sampler = _tiny_sampler(model)
    steps = [
        sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
        for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
    ]
    expected = [sampler.forward(step) for step in steps]

    torch._dynamo.reset()
    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
    first = sampler.forward(steps[0])
    graphs = torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before
    with torch._dynamo.config.patch(error_on_recompile=True):
        second = sampler.forward(steps[1])

    for (video, audio), (video_ref, audio_ref) in zip((first, second), expected, strict=True):
        assert torch.equal(video, video_ref)
        assert torch.equal(audio, audio_ref)
    # pre adapter, post adapter, three MoE-layer regions, two dense-layer regions
    assert graphs == 7


def test_pipeline_setup_compile_forwards_config_dynamic() -> None:
    from vllm_omni.diffusion.data import OmniDiffusionConfig
    from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline

    calls: list[dict[str, object]] = []

    class _Transformer(torch.nn.Module):
        def compile_regions(self, **compile_kwargs: object) -> None:
            calls.append(compile_kwargs)

    pipeline = object.__new__(Magi2Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.od_config = OmniDiffusionConfig(
        model="sand-ai/MAGI-2-preview",
        model_class_name="Magi2Pipeline",
        diffusion_compile_dynamic=False,
    )
    pipeline.transformer = _Transformer()

    pipeline.setup_compile()

    assert calls == [{"fullgraph": True, "dynamic": False, "options": {"emulate_precision_casts": True}}]


def test_pipeline_setup_compile_warns_when_granularity_is_not_regional(caplog: pytest.LogCaptureFixture) -> None:
    from vllm_omni.diffusion.data import OmniDiffusionConfig
    from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline

    calls: list[dict[str, object]] = []

    class _Transformer(torch.nn.Module):
        def compile_regions(self, **compile_kwargs: object) -> None:
            calls.append(compile_kwargs)

    pipeline = object.__new__(Magi2Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.od_config = OmniDiffusionConfig(
        model="sand-ai/MAGI-2-preview",
        model_class_name="Magi2Pipeline",
        diffusion_compile_granularity="full",
    )
    pipeline.transformer = _Transformer()

    pipeline.setup_compile()

    assert calls == [{"fullgraph": True, "dynamic": True, "options": {"emulate_precision_casts": True}}]
    assert "diffusion_compile_granularity='full'" in caplog.text


def test_compiled_graphs_are_shared_across_layers_of_one_kind() -> None:
    model = _tiny_model(moe_layers=2)
    sampler = _tiny_sampler(model)
    step = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([900.0]), cfg_config=CFGConfig())
    expected = sampler.forward(step)

    torch._dynamo.reset()
    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
    actual = sampler.forward(step)

    assert torch.equal(actual[0], expected[0])
    # Two MoE layers and two dense layers still compile one graph per region
    # kind: pre adapter, post adapter, three MoE regions, two dense regions.
    assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7


def test_compiled_regions_survive_new_requests_without_recompiling() -> None:
    model = _tiny_model()
    sampler = _tiny_sampler(model)
    t = torch.tensor([900.0])
    first_request = sampler.prepare_model_input(**_sampler_tensors(0), t=t, cfg_config=CFGConfig())
    second_request = sampler.prepare_model_input(**_sampler_tensors(1), t=t, cfg_config=CFGConfig())
    longer_request = sampler.prepare_model_input(**_longer_text_tensors(2), t=t, cfg_config=CFGConfig())
    expected = [sampler.forward(request) for request in (first_request, second_request, longer_request)]

    torch._dynamo.reset()
    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
    actual = [sampler.forward(first_request)]
    with torch._dynamo.config.patch(error_on_recompile=True):
        actual.append(sampler.forward(second_request))
        actual.append(sampler.forward(longer_request))

    for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
        assert torch.equal(video, video_ref)
        assert torch.equal(audio, audio_ref)
    assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7