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

from collections import defaultdict

import pytest
import torch
import torch.nn as nn
from diffusers.loaders.lora_conversion_utils import (
    _convert_non_diffusers_qwen_lora_to_diffusers,
)
from pytest_mock import MockerFixture
from torch.testing import assert_close

from vllm_omni.diffusion.lora.loader import (
    LoraLoaderMixin,
    QwenImageLoraLoaderMixin,
    WanLoraLoaderMixin,
    _prepare_lora_delta,
    _remap_state_dict_keys,
    _resolve_lora_compute_device,
)

HEAD_DIM = 24
RANK = 2
NUM_LAYERS = 10

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


class DummyTransformerBlock(nn.Module):
    def __init__(
        self,
        d_model: int,
        bias: bool = False,
        stacked_params: bool = False,
    ):
        super().__init__()
        self.attn = nn.Module()
        if stacked_params:
            self.attn.to_qkv = nn.Linear(d_model * 3, d_model, bias=bias)
        else:
            self.attn.to_q = nn.Linear(d_model, d_model, bias=bias)
            self.attn.to_k = nn.Linear(d_model, d_model, bias=bias)
            self.attn.to_v = nn.Linear(d_model, d_model, bias=bias)
        self.attn.to_out = nn.Linear(d_model, d_model, bias=bias)


class DummyTransformer(nn.Module):
    def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
        super().__init__()

        self.blocks = nn.ModuleList(
            [DummyTransformerBlock(d_model, bias=bias, stacked_params=stacked_params) for _ in range(num_layers)]
        )
        if stacked_params:
            self.stacked_params_mapping = [
                (".attn.to_qkv", ".attn.to_q", "q"),
                (".attn.to_qkv", ".attn.to_k", "k"),
                (".attn.to_qkv", ".attn.to_v", "v"),
            ]


class DummyPipeline(nn.Module, LoraLoaderMixin):
    def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
        super().__init__()
        self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=stacked_params)


# ======================================================
# helper functions for test
# ======================================================
def _filter_parameters(named_params, mode: str = "all"):
    if mode == "all":
        return named_params

    filtered = []
    for name, param in named_params:
        if name.endswith(mode):
            filtered.append((name, param))

    return filtered


def assert_params_equal(module, original_params, mode: str = "all"):
    for name, param in _filter_parameters(module.named_parameters(), mode):
        assert_close(param, original_params[name])


def assert_params_not_equal(module, original_params, mode: str = "all"):
    for name, param in _filter_parameters(module.named_parameters(), mode):
        assert not torch.allclose(param, original_params[name])


def make_dummy_lora_state_dict(
    base_key,
    weight_a: torch.Tensor | None = None,
    weight_b: torch.Tensor | None = None,
    weight_bias: torch.Tensor | None = None,
    bias: bool = False,
    lora_a_suffix: str = "lora_A.weight",
    lora_b_suffix: str = "lora_B.weight",
    lora_bias_suffix: str = "bias",
):
    """
    Create a dummy LoRA state dict.

    Returns:
        tuple: (state_dict, expected_delta)
    """
    key_a = f"{base_key}.{lora_a_suffix}"
    key_b = f"{base_key}.{lora_b_suffix}"
    key_bias = f"{base_key}.{lora_bias_suffix}"

    if weight_a is None:
        weight_a = torch.randn(RANK, HEAD_DIM)
    if weight_b is None:
        weight_b = torch.randn(HEAD_DIM, RANK)

    state_dict = {}
    expected_delta = None

    if bias:
        state_dict[key_a] = weight_a
        state_dict[key_b] = weight_b
        weight_bias = torch.randn(HEAD_DIM)
        state_dict[key_bias] = weight_bias
        expected_delta = weight_bias
    else:
        state_dict[key_a] = weight_a
        state_dict[key_b] = weight_b
        expected_delta = weight_b @ weight_a

    return state_dict, expected_delta


def make_lora_state_dict_for_module(
    module,
    lora_rank=RANK,
    bias=False,
    lora_a_suffix: str = "lora_A.weight",
    lora_b_suffix: str = "lora_B.weight",
    lora_bias_suffix: str = "bias",
    prefix: str = "transformer",
    remap_proj_out: bool = True,
):
    param_to_weight_names = defaultdict(list)
    if hasattr(module, "stacked_params_mapping"):
        for param_name, weight_name, _ in module.stacked_params_mapping:
            param_to_weight_names[param_name].append(weight_name)

    state_dict = {}
    for name, param in module.named_parameters(prefix=prefix):
        if name.endswith(".weight"):
            base_key = name[: -len(".weight")]
            if remap_proj_out and base_key.endswith(".to_out"):
                base_key = base_key.replace(".to_out", ".to_out.0")

            d_in, d_out = param.shape
            is_stacked_param = False
            for param_name, weight_names in param_to_weight_names.items():
                if param_name not in name:
                    continue
                is_stacked_param = True
                target_in = d_in // len(weight_names)
                for weight_name in weight_names:
                    lora_key = base_key.replace(param_name, weight_name)
                    weight_a = torch.randn(lora_rank, d_out)
                    weight_b = torch.randn(target_in, lora_rank)
                    state_dict[f"{lora_key}.{lora_a_suffix}"] = weight_a
                    state_dict[f"{lora_key}.{lora_b_suffix}"] = weight_b

            if is_stacked_param:
                continue

            weight_a = torch.randn(lora_rank, d_out)
            weight_b = torch.randn(d_in, lora_rank)
            key_a = f"{base_key}.{lora_a_suffix}"
            key_b = f"{base_key}.{lora_b_suffix}"
            state_dict[key_a] = weight_a
            state_dict[key_b] = weight_b
        elif name.endswith(".bias") and bias:
            base_key = name[: -len(".bias")]
            if remap_proj_out and base_key.endswith(".to_out"):
                base_key = base_key.replace(".to_out", ".to_out.0")

            d_bias = param.shape[0]
            is_stacked_param = False
            for param_name, weight_names in param_to_weight_names.items():
                if param_name not in name:
                    continue
                is_stacked_param = True
                target_bias = d_bias // len(weight_names)
                for weight_name in weight_names:
                    lora_key = base_key.replace(param_name, weight_name)
                    weight_bias = torch.randn(target_bias)
                    state_dict[f"{lora_key}.{lora_bias_suffix}"] = weight_bias

            if is_stacked_param:
                continue

            key_bias = f"{base_key}.{lora_bias_suffix}"
            state_dict[key_bias] = torch.randn_like(param)

    return state_dict


# ======================================================
# Test loader helper functions
# ======================================================
def test_prepare_lora_delta_with_stacked_params():
    base_key = "transformer.attn.to_qkv"

    sd0, delta0 = make_dummy_lora_state_dict("transformer.attn.to_q")
    sd1, delta1 = make_dummy_lora_state_dict("transformer.attn.to_k")
    sd2, delta2 = make_dummy_lora_state_dict("transformer.attn.to_v")

    lora_state_dict = {}
    lora_state_dict.update(sd0)
    lora_state_dict.update(sd1)
    lora_state_dict.update(sd2)
    expected_delta = torch.concat([delta0, delta1, delta2])

    stacked_mapping = {".attn.to_qkv": [".attn.to_q", ".attn.to_k", ".attn.to_v"]}

    actual_delta, used_keys = _prepare_lora_delta(
        lora_state_dict,
        base_key,
        stacked_mapping,
    )

    # validate the delta
    assert actual_delta is not None
    assert_close(actual_delta, expected_delta)

    # validate the used keys
    key_q_a = "transformer.attn.to_q.lora_A.weight"
    key_q_b = "transformer.attn.to_q.lora_B.weight"
    key_k_a = "transformer.attn.to_k.lora_A.weight"
    key_k_b = "transformer.attn.to_k.lora_B.weight"
    key_v_a = "transformer.attn.to_v.lora_A.weight"
    key_v_b = "transformer.attn.to_v.lora_B.weight"

    for key in [key_q_a, key_q_b, key_k_a, key_k_b, key_v_a, key_v_b]:
        assert key in used_keys


def test_remap_state_dict_keys_matched_two():
    lora_state_dict, _ = make_dummy_lora_state_dict(
        "transformer.attn.to_out.0", bias=True, lora_bias_suffix="lora_B.bias"
    )

    lora_state_dict = _remap_state_dict_keys(
        lora_state_dict,
        [(".to_out.0.", ".to_out."), (".lora_B.bias", ".bias")],
    )

    assert len(lora_state_dict) == 3
    assert "transformer.attn.to_out.bias" in lora_state_dict


# ======================================================
# Test LoraLoaderMixin
# ======================================================
class TestLoraLoaderMixin:
    @pytest.mark.parametrize(
        "bias, stacked_params",
        [
            (False, False),
            (True, False),
            (False, True),
            (True, True),
        ],
        ids=[
            "basic",
            "bias",
            "stacked_params",
            "bias_stacked_params",
        ],
    )
    def test_load_lora_into_module(self, bias: bool, stacked_params: bool):
        pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM, bias=bias, stacked_params=stacked_params)
        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}

        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer, bias=bias)
        assert len(lora_state_dict) > 0

        # remap the keys to vllm-omni format
        lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])

        used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
        # validate the parameters are updated
        assert_params_not_equal(pipeline.transformer, original_params)

        # validate all keys are used
        missing_keys = set(lora_state_dict.keys()) - used_keys
        assert len(missing_keys) == 0

    def test_unload_module_lora(self):
        pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}

        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
        # validate the state dict is valid
        assert len(lora_state_dict) > 0

        # remap the keys to vllm-omni format
        lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])

        used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)

        # validate the weights are updated
        assert_params_not_equal(pipeline.transformer, original_params)

        # validate all keys are used
        missing_keys = set(lora_state_dict.keys()) - used_keys
        assert len(missing_keys) == 0

        unload_used_keys = pipeline.unload_module_lora(lora_state_dict, pipeline.transformer)
        missing_keys = set(lora_state_dict.keys()) - unload_used_keys
        assert len(missing_keys) == 0
        assert len(unload_used_keys - used_keys) == 0

        # validate the weights are restored
        assert_params_equal(pipeline.transformer, original_params)


# ======================================================
# Test QwenImageLoraLoaderMixin
# ======================================================
class DummyQwenImagePipeline(nn.Module, QwenImageLoraLoaderMixin):
    def __init__(self, num_layers: int, d_model: int, bias: bool = False):
        super().__init__()
        self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=True)


class TestQwenImageLoraLoaderMixin:
    def test_load_lora_weights(self, mocker: MockerFixture):
        pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}

        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
        assert len(lora_state_dict) > 0

        mocker.patch(
            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
        )
        pipeline.load_lora_weights(lora_state_dict, "adapter0")

        # validate the weights are updated after lora loaded
        assert_params_not_equal(pipeline.transformer, original_params)

        # validate lora_loaded map is updated after lora loaded
        assert "adapter0" in pipeline.lora_loaded

    def test_unload_lora_weights(self, mocker: MockerFixture):
        pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}

        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
        assert len(lora_state_dict) > 0

        mocker.patch(
            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
        )
        pipeline.load_lora_weights(lora_state_dict, "adapter0")
        # validate the weights are updated after lora loaded
        assert_params_not_equal(pipeline.transformer, original_params)

        # validate lora_loaded map is updated after lora loaded
        assert "adapter0" in pipeline.lora_loaded

        pipeline.unload_lora_weights("adapter0")
        # validate the weights are restored after lora unloaded
        assert_params_equal(pipeline.transformer, original_params)

        # validate lora_loaded map is updated
        assert "adapter0" not in pipeline.lora_loaded

    def test_load_diffusers_lora_applies_non_unit_alpha(self, mocker: MockerFixture):
        pipeline = DummyQwenImagePipeline(1, HEAD_DIM)
        to_out = pipeline.transformer.blocks[0].attn.to_out
        original_weight = to_out.weight.detach().clone()
        base_key = "blocks.0.attn.to_out.0"
        lora_a = torch.randn(RANK, HEAD_DIM)
        lora_b = torch.randn(HEAD_DIM, RANK)
        alpha = torch.tensor(6.0)
        state_dict = {
            f"{base_key}.lora_A.weight": lora_a,
            f"{base_key}.lora_B.weight": lora_b,
            f"{base_key}.alpha": alpha,
        }
        mocker.patch(
            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
        )

        pipeline.load_lora_weights(state_dict, "adapter0")

        expected_weight = original_weight + (lora_b @ lora_a) * (alpha.item() / RANK)
        assert_close(to_out.weight, expected_weight)

        pipeline.unload_lora_weights("adapter0")
        assert_close(to_out.weight, original_weight)


# ======================================================
# Test WanLoraLoaderMixin
# ======================================================
class DummyWanPipeline(nn.Module, WanLoraLoaderMixin):
    def __init__(self, num_layers: int, d_model: int):
        super().__init__()
        self.has_transformer_2 = True
        self.transformer = DummyTransformer(num_layers, d_model, stacked_params=True)
        self.transformer_2 = DummyTransformer(num_layers, d_model, stacked_params=True)


def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
    pipeline = DummyWanPipeline(NUM_LAYERS, HEAD_DIM)
    original_high = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
    original_low = {name: param.clone() for name, param in pipeline.transformer_2.named_parameters()}
    high_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
    low_state_dict = make_lora_state_dict_for_module(pipeline.transformer_2)

    mocker.patch(
        "vllm_omni.diffusion.lora.loader._load_lora_state_dict",
        side_effect=[high_state_dict, low_state_dict],
    )

    pipeline.load_lora_weights(["high.safetensors", "low.safetensors"], "adapter0")

    assert_params_not_equal(pipeline.transformer, original_high)
    assert_params_not_equal(pipeline.transformer_2, original_low)
    assert set(pipeline.lora_loaded["adapter0"]) == {"transformer", "transformer_2"}

    pipeline.unload_lora_weights("adapter0")

    assert_params_equal(pipeline.transformer, original_high)
    assert_params_equal(pipeline.transformer_2, original_low)
    assert "adapter0" not in pipeline.lora_loaded


def test_prepare_lora_delta_disambiguates_prefix_overlapping_stacked_names():
    """A fused name can be a prefix of another fused name.

    SenseNova-U1 carries both `.qkv_proj` (understanding) and `.qkv_proj_mot_gen`
    (generation). Selecting stacked mappings by substring matched both rules for
    the generation weight and concatenated six shards into a delta twice as tall
    as the parameter.
    """
    param_to_weight_names = {
        ".qkv_proj_mot_gen": [".q_proj_mot_gen", ".k_proj_mot_gen", ".v_proj_mot_gen"],
        ".qkv_proj": [".q_proj", ".k_proj", ".v_proj"],
    }
    state_dict = {}
    for part in ("q", "k", "v"):
        for tower in ("_mot_gen", ""):
            base = f"layers.0.self_attn.{part}_proj{tower}"
            state_dict[f"{base}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
            state_dict[f"{base}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)

    for base_key in ("layers.0.self_attn.qkv_proj_mot_gen", "layers.0.self_attn.qkv_proj"):
        delta, _ = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
        assert delta.shape == (3 * HEAD_DIM, HEAD_DIM), base_key


@pytest.mark.parametrize(
    "fused, shards",
    [
        # qwen_image_transformer.py stacked_params_mapping
        (".to_qkv", (".to_q", ".to_k", ".to_v")),
        (".add_kv_proj", (".add_q_proj", ".add_k_proj", ".add_v_proj")),
        # wan2_2_transformer.py stacked_params_mapping
        (".attn1.to_qkv", (".attn1.to_q", ".attn1.to_k", ".attn1.to_v")),
        # wan2_2_s2v_transformer.py stacked_params_mapping
        (".self_attn.to_qkv", (".self_attn.q", ".self_attn.k", ".self_attn.v")),
    ],
)
def test_prepare_lora_delta_unchanged_for_existing_fused_layouts(fused, shards):
    """None of the fused names shipped today overlap by prefix, so matching on the
    tail must select exactly the shards substring matching selected."""
    param_to_weight_names = {fused: list(shards)}
    base_key = f"blocks.0{fused}"
    state_dict = {}
    for shard in shards:
        state_dict[f"blocks.0{shard}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
        state_dict[f"blocks.0{shard}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)

    delta, used_keys = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
    assert delta.shape == (len(shards) * HEAD_DIM, HEAD_DIM)
    assert len(used_keys) == 2 * len(shards)


def test_lora_loader_mixin_lora_is_fused_property(mocker: MockerFixture):
    pipeline = DummyQwenImagePipeline(1, HEAD_DIM)
    assert pipeline.lora_is_fused is False

    lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
    mocker.patch(
        "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
        return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
    )
    pipeline.load_lora_weights(lora_state_dict, "adapter0")
    assert pipeline.lora_is_fused is True

    # Test explicit setter
    pipeline2 = DummyQwenImagePipeline(1, HEAD_DIM)
    assert pipeline2.lora_is_fused is False
    pipeline2.lora_is_fused = True
    assert pipeline2.lora_is_fused is True


def test_prepare_lora_delta_with_compute_device():
    base_key = "blocks.0.attn.to_q"
    state_dict = {
        f"{base_key}.lora_A.weight": torch.ones(RANK, HEAD_DIM),
        f"{base_key}.lora_B.weight": torch.ones(HEAD_DIM, RANK),
    }
    delta, used_keys = _prepare_lora_delta(state_dict, base_key, compute_device=torch.device("cpu"))
    assert delta is not None
    assert delta.shape == (HEAD_DIM, HEAD_DIM)
    assert len(used_keys) == 2


def test_resolve_lora_compute_device_probes_unusable_accelerator(mocker: MockerFixture):
    """cuda:N/npu:N may be reported without a physical device; the probe must reject it."""
    mocker.patch(
        "vllm_omni.diffusion.lora.loader.get_local_device",
        return_value=torch.device("cuda:999"),
    )
    assert _resolve_lora_compute_device() is None


def test_load_lora_into_module_cpu_weights_match_cpu_reference(mocker: MockerFixture):
    """Regression test: params on CPU used to force compute_device=None and skip the
    #7005 accelerator path. With an unusable accelerator, fusion must fall back to
    CPU and still produce base + lora_B @ lora_A."""
    pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
    original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}

    lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
    lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])

    mocker.patch(
        "vllm_omni.diffusion.lora.loader.get_local_device",
        return_value=torch.device("cuda:999"),
    )

    used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
    assert len(used_keys) > 0
    assert not (set(lora_state_dict.keys()) - used_keys)

    for name, param in pipeline.transformer.named_parameters():
        if not name.endswith(".weight"):
            continue
        base_key = name[: -len(".weight")]
        key_a = f"{base_key}.lora_A.weight"
        key_b = f"{base_key}.lora_B.weight"
        if key_a in lora_state_dict and key_b in lora_state_dict:
            expected = original_params[name] + (
                lora_state_dict[key_b].to(param.dtype) @ lora_state_dict[key_a].to(param.dtype)
            )
            assert_close(param, expected)


def test_load_lora_into_module_forwards_usable_accel_device(mocker: MockerFixture):
    """A usable accelerator must be forwarded as compute_device for CPU params.
    A spy forces the actual math back onto CPU so the test runs without a device."""
    fake_dev = torch.device("cuda:0")
    mocker.patch(
        "vllm_omni.diffusion.lora.loader._resolve_lora_compute_device",
        return_value=fake_dev,
    )

    seen_devices = []
    original = _prepare_lora_delta

    def spy(state_dict, base_key, *args, **kwargs):
        seen_devices.append(kwargs.get("compute_device"))
        kwargs["compute_device"] = None
        return original(state_dict, base_key, *args, **kwargs)

    mocker.patch(
        "vllm_omni.diffusion.lora.loader._prepare_lora_delta",
        side_effect=spy,
    )

    pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
    original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}

    lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
    lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])

    used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)

    # Every CPU param with a LoRA delta must have received the accelerator device.
    assert seen_devices
    assert all(d == fake_dev for d in seen_devices)
    assert not (set(lora_state_dict.keys()) - used_keys)
    assert_params_not_equal(pipeline.transformer, original_params)