"""Whole-parameter LoRA payload: key classification, application, and reporting.

Nothing here loads a model. Most tests exercise key algebra and arithmetic on the CPU;
one regression test verifies cross-device adapter loading when CUDA is available.
"""
import logging

import pytest
import torch
from safetensors.torch import save_file

from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
from fastvideo.models.loader import fsdp_load
from fastvideo.models.loader.lora_patch import DenseLoRAPatch, normalize_lora_key
from fastvideo.models.loader.utils import get_param_names_mapping


def write_adapter(tmp_path, tensors, name="adapter_model.safetensors"):
    path = tmp_path / name
    save_file(tensors, str(path))
    return str(path)


# --------------------------------------------------------------------------- keys


@pytest.mark.parametrize(
    "raw, expected",
    [
        # PEFT, the spelling FastVideo's own extractor emits.
        ("blocks.0.attn.to_q.lora_A.weight", "blocks.0.attn.to_q.lora_A.weight"),
        # PEFT with a named adapter: the name sits between the factor and `.weight`.
        ("blocks.0.attn.to_q.lora_A.default.weight", "blocks.0.attn.to_q.lora_A.weight"),
        ("blocks.0.attn.to_q.lora_B.turbo_v2.weight", "blocks.0.attn.to_q.lora_B.weight"),
        # kohya / ComfyUI up-down spelling, with and without `.weight`.
        ("blocks.0.attn.to_q.lora_down.weight", "blocks.0.attn.to_q.lora_A.weight"),
        ("blocks.0.attn.to_q.lora_up.weight", "blocks.0.attn.to_q.lora_B.weight"),
        ("blocks.0.attn.to_q.lora_down", "blocks.0.attn.to_q.lora_A"),
        # kohya alpha.
        ("blocks.0.attn.to_q.alpha", "blocks.0.attn.to_q.lora_alpha"),
        # The `diffusion_model.` prefix ComfyUI-format files carry.
        ("diffusion_model.blocks.0.mlp.fc1.lora_up.weight", "blocks.0.mlp.fc1.lora_B.weight"),
    ],
)
def test_normalize_lora_key_accepts_every_published_spelling(raw, expected):
    assert normalize_lora_key(raw) == expected


@pytest.mark.parametrize("raw", [
    "blocks.0.norm1.diff",
    "blocks.0.attn.to_q.diff_b",
    "blocks.0.attn.to_gate_compress.set_weight",
    "blocks.0.scale_shift_table.diff_param",
    "blocks.0.extra_table.set_param",
    "blocks.0.attn.to_q.dora_scale",
])
def test_normalize_lora_key_disclaims_non_low_rank_keys(raw):
    """These belong to the loader-side patch, not the module merge path."""
    assert normalize_lora_key(raw) is None


def test_set_weight_is_not_mangled_into_set():
    """`.replace('.weight', '')` downstream would turn `.set_weight` into `.set`."""
    assert normalize_lora_key("blocks.0.attn.to_gate_compress.set_weight") is None


# --------------------------------------------------------------------- classification


def test_from_adapter_returns_none_for_a_purely_low_rank_adapter(tmp_path):
    path = write_adapter(
        tmp_path, {
            "blocks.0.attn.to_q.lora_A.weight": torch.zeros(4, 8),
            "blocks.0.attn.to_q.lora_B.weight": torch.zeros(8, 4),
        })
    assert DenseLoRAPatch.from_adapter(path) is None


def test_from_adapter_returns_none_without_a_path():
    assert DenseLoRAPatch.from_adapter(None) is None


def test_from_adapter_splits_additive_from_replacement(tmp_path):
    path = write_adapter(
        tmp_path, {
            "blocks.0.attn.to_q.lora_A.weight": torch.zeros(4, 8),
            "blocks.0.norm1.diff": torch.zeros(8),
            "blocks.0.attn.to_q.diff_b": torch.zeros(8),
            "blocks.0.scale_shift_table.diff_param": torch.zeros(8),
            "blocks.0.attn.to_gate_compress.set_weight": torch.zeros(8, 8),
            "blocks.0.extra_table.set_param": torch.zeros(8),
        })
    patch = DenseLoRAPatch.from_adapter(path)
    assert patch is not None
    # Weight/bias suffixes restore their parameter suffix; generic parameter
    # payloads preserve the complete name.
    assert set(patch._additive) == {
        "blocks.0.norm1.weight",
        "blocks.0.attn.to_q.bias",
        "blocks.0.scale_shift_table",
    }
    assert set(patch._replacement) == {
        "blocks.0.attn.to_gate_compress.weight",
        "blocks.0.extra_table",
    }


def test_dense_suffix_is_not_hidden_by_an_alpha_like_parameter_name(tmp_path):
    path = write_adapter(tmp_path, {"blocks.0.alpha.diff_param": torch.zeros(4)})
    patch = DenseLoRAPatch.from_adapter(path)
    assert set(patch._additive) == {"blocks.0.alpha"}


def test_param_names_mapping_is_applied_to_dense_keys(tmp_path):
    """An adapter written against the published layout needs no separate rename table."""
    path = write_adapter(tmp_path, {"blocks.0.ff.net.0.proj.diff": torch.zeros(4)})

    def mapping(name):
        return name.replace("ff.net.0.proj", "ff.fc_in"), None, None

    patch = DenseLoRAPatch.from_adapter(path, mapping)
    assert set(patch._additive) == {"blocks.0.ff.fc_in.weight"}


def test_official_wan_mixed_adapter_maps_dense_key_like_factors(tmp_path):
    path = write_adapter(
        tmp_path, {
            "blocks.0.self_attn.q.lora_A.weight": torch.zeros(2, 4),
            "blocks.0.self_attn.q.lora_B.weight": torch.zeros(4, 2),
            "blocks.0.self_attn.k.diff": torch.zeros(4),
        })
    config = WanVideoConfig()
    patch = DenseLoRAPatch.from_adapter(
        path,
        get_param_names_mapping(config.param_names_mapping),
        lora_param_names_mapping=get_param_names_mapping(config.lora_param_names_mapping),
    )
    assert set(patch._additive) == {"blocks.0.to_k.weight"}


def _stub_fsdp_loading(monkeypatch):
    monkeypatch.setattr(fsdp_load, "set_mixed_precision_policy", lambda **_: None)
    monkeypatch.setattr(fsdp_load, "safetensors_weights_iterator", lambda *_args, **_kwargs: iter(()))
    monkeypatch.setattr(fsdp_load, "load_model_from_full_model_state_dict", lambda *_args, **_kwargs: None)
    monkeypatch.setattr(fsdp_load, "_maybe_quantize_model", lambda _model, **_kwargs: None)


def _load_tiny_model(model_cls, *, lora_path=None, lora_strength=1.0, init_params=None):
    return fsdp_load.maybe_load_fsdp_model(
        model_cls=model_cls,
        init_params=init_params or {},
        weight_dir_list=[],
        device=torch.device("cpu"),
        hsdp_replicate_dim=1,
        hsdp_shard_dim=1,
        default_dtype=torch.float32,
        param_dtype=torch.float32,
        reduce_dtype=torch.float32,
        training_mode=False,
        pin_cpu_memory=False,
        lora_path=lora_path,
        lora_strength=lora_strength,
    )


def test_fsdp_loader_forwards_lora_specific_mapping(monkeypatch):
    captured = {}

    class TinyModel(torch.nn.Module):
        param_names_mapping = {r"^hf\.(.*)$": r"custom.\1"}
        lora_param_names_mapping = {r"^official\.(.*)$": r"hf.\1"}

    def capture_patch(cls,
                      lora_path,
                      param_names_mapping,
                      *,
                      lora_param_names_mapping=None,
                      strength=1.0):
        captured["path"] = lora_path
        captured["regular"] = param_names_mapping("hf.weight")[0]
        captured["lora"] = lora_param_names_mapping("official.weight")[0]
        captured["strength"] = strength
        return None

    monkeypatch.setattr(DenseLoRAPatch, "from_adapter", classmethod(capture_patch))
    _stub_fsdp_loading(monkeypatch)

    _load_tiny_model(TinyModel, lora_path="adapter.safetensors", lora_strength=0.75)

    assert captured == {
        "path": "adapter.safetensors",
        "regular": "custom.weight",
        "lora": "hf.weight",
        "strength": 0.75,
    }


def test_fsdp_loader_does_not_consult_lora_mapping_without_adapter(monkeypatch):
    class TinyModel(torch.nn.Module):
        param_names_mapping = {r"^hf\.(.*)$": r"custom.\1"}

    def unexpected_patch(*_args, **_kwargs):
        pytest.fail("the no-LoRA load path must not inspect an adapter")

    monkeypatch.setattr(DenseLoRAPatch, "from_adapter", classmethod(unexpected_patch))
    _stub_fsdp_loading(monkeypatch)

    model = _load_tiny_model(TinyModel)

    assert isinstance(model, TinyModel)
    assert not hasattr(model, "lora_param_names_mapping")


def test_fsdp_loader_allows_adapter_without_lora_specific_mapping(monkeypatch):
    captured = {}

    class TinyModel(torch.nn.Module):
        param_names_mapping = {r"^hf\.(.*)$": r"custom.\1"}

    def capture_patch(cls,
                      lora_path,
                      param_names_mapping,
                      *,
                      lora_param_names_mapping=None,
                      strength=1.0):
        captured["path"] = lora_path
        captured["regular"] = param_names_mapping("hf.weight")[0]
        captured["lora"] = lora_param_names_mapping
        captured["strength"] = strength
        return None

    monkeypatch.setattr(DenseLoRAPatch, "from_adapter", classmethod(capture_patch))
    _stub_fsdp_loading(monkeypatch)

    _load_tiny_model(TinyModel, lora_path="adapter.safetensors")

    assert captured == {
        "path": "adapter.safetensors",
        "regular": "custom.weight",
        "lora": None,
        "strength": 1.0,
    }


def test_fused_target_is_refused_rather_than_guessed(tmp_path):
    path = write_adapter(tmp_path, {"blocks.0.attn.to_q.diff": torch.zeros(4)})

    def fusing_mapping(name):
        return "blocks.0.attn.to_qkv.weight", 0, 3

    with pytest.raises(NotImplementedError, match="fused parameter"):
        DenseLoRAPatch.from_adapter(path, fusing_mapping)


def test_two_keys_onto_one_parameter_is_an_error(tmp_path):
    path = write_adapter(tmp_path, {
        "blocks.0.norm1.diff": torch.zeros(4),
        "blocks.0.norm1.set_weight": torch.zeros(4),
    })
    # `.diff` and `.set_weight` land in different tables, so they do not collide; two
    # keys of the *same* kind must.
    patch = DenseLoRAPatch.from_adapter(path)
    assert set(patch._additive) == {"blocks.0.norm1.weight"}
    assert set(patch._replacement) == {"blocks.0.norm1.weight"}


# ---------------------------------------------------------------------- application


def test_apply_to_adds_the_delta(tmp_path):
    path = write_adapter(tmp_path, {"blocks.0.norm1.diff": torch.full((4, ), 0.25)})
    patch = DenseLoRAPatch.from_adapter(path)
    out = patch.apply_to("blocks.0.norm1.weight", torch.ones(4))
    assert torch.allclose(out, torch.full((4, ), 1.25))


def test_apply_to_supports_standalone_parameters(tmp_path):
    path = write_adapter(tmp_path, {"blocks.0.scale_shift_table.diff_param": torch.full((4, ), 0.25)})
    patch = DenseLoRAPatch.from_adapter(path)
    out = patch.apply_to("blocks.0.scale_shift_table", torch.ones(4))
    assert torch.allclose(out, torch.full((4, ), 1.25))


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for cross-device adapter loading")
def test_apply_to_moves_cpu_delta_to_base_tensor_device(tmp_path):
    """A CPU safetensors delta follows a directly loaded CUDA checkpoint tensor."""
    path = write_adapter(tmp_path, {"blocks.0.norm1.diff": torch.full((4, ), 0.25)})
    patch = DenseLoRAPatch.from_adapter(path)

    out = patch.apply_to("blocks.0.norm1.weight", torch.ones(4, device="cuda"))

    assert out.device.type == "cuda"
    assert torch.allclose(out.cpu(), torch.full((4, ), 1.25))


def test_apply_to_leaves_unrelated_parameters_untouched(tmp_path):
    path = write_adapter(tmp_path, {"blocks.0.norm1.diff": torch.full((4, ), 0.25)})
    patch = DenseLoRAPatch.from_adapter(path)
    base = torch.ones(4)
    assert patch.apply_to("blocks.9.norm1.weight", base) is base


def test_apply_to_sums_in_float32_so_a_small_delta_survives(tmp_path):
    """A bfloat16 add would quantize away a delta far below the base weight.

    The VSA gates and the norm deltas both sit three to four orders of magnitude under
    the weights they modify, which is exactly where bfloat16's 8-bit significand drops
    the change entirely.
    """
    delta = 1e-4
    path = write_adapter(tmp_path, {"blocks.0.norm1.diff": torch.full((4, ), delta, dtype=torch.bfloat16)})
    patch = DenseLoRAPatch.from_adapter(path)
    base = torch.ones(4, dtype=torch.bfloat16)

    out = patch.apply_to("blocks.0.norm1.weight", base)
    assert out.dtype == torch.float32
    assert (out - 1.0).abs().min() > 0, "delta was rounded away"

    naive = (base + torch.full((4, ), delta, dtype=torch.bfloat16))
    assert torch.equal(naive, base), "precondition: the bfloat16 add is what loses it"


def test_apply_to_rejects_a_shape_mismatch(tmp_path):
    path = write_adapter(tmp_path, {"blocks.0.norm1.diff": torch.zeros(8)})
    patch = DenseLoRAPatch.from_adapter(path)
    with pytest.raises(ValueError, match="shape"):
        patch.apply_to("blocks.0.norm1.weight", torch.zeros(4))


def test_replacement_for_supports_standalone_parameters(tmp_path):
    table = torch.randn(6, 4)
    path = write_adapter(tmp_path, {"blocks.0.extra_table.set_param": table})
    patch = DenseLoRAPatch.from_adapter(path)
    assert torch.equal(patch.replacement_for("blocks.0.extra_table"), table)


def test_replacement_for_returns_the_whole_tensor(tmp_path):
    gate = torch.randn(6, 4)
    path = write_adapter(tmp_path, {"blocks.0.attn.to_gate_compress.set_weight": gate})
    patch = DenseLoRAPatch.from_adapter(path)
    assert torch.equal(patch.replacement_for("blocks.0.attn.to_gate_compress.weight"), gate)
    assert patch.replacement_for("blocks.0.attn.to_q.weight") is None


# ------------------------------------------------------------------------ reporting


def test_unapplied_keys_are_named_in_the_log(tmp_path, caplog):
    path = write_adapter(tmp_path, {
        "blocks.0.norm1.diff": torch.zeros(4),
        "blocks.77.norm1.diff": torch.zeros(4),
    })
    patch = DenseLoRAPatch.from_adapter(path)
    patch.apply_to("blocks.0.norm1.weight", torch.zeros(4))

    with caplog.at_level(logging.WARNING):
        patch.report_unapplied()

    assert "blocks.77.norm1.diff" in caplog.text
    assert "blocks.0.norm1.diff" not in caplog.text


def test_a_fully_applied_payload_warns_about_nothing(tmp_path, caplog):
    path = write_adapter(tmp_path, {"blocks.0.norm1.diff": torch.zeros(4)})
    patch = DenseLoRAPatch.from_adapter(path)
    patch.apply_to("blocks.0.norm1.weight", torch.zeros(4))

    with caplog.at_level(logging.WARNING):
        patch.report_unapplied()

    assert [r for r in caplog.records if r.levelno >= logging.WARNING] == []


def _stub_packed_nvfp4_export(monkeypatch, tmp_path):
    from fastvideo.layers.quantization import nvfp4_config

    packed_path = tmp_path / "packed.safetensors"
    save_file({"placeholder": torch.zeros(1)}, str(packed_path))
    loaded = []
    monkeypatch.setattr(nvfp4_config, "find_minimax_h3_nvfp4_dit_export", lambda _paths: str(packed_path))
    monkeypatch.setattr(nvfp4_config, "nvfp4_linear_weight_param_names", lambda _model: {"linear.weight"})
    monkeypatch.setattr(nvfp4_config, "load_minimax_h3_nvfp4_dit_export",
                        lambda _model, path, device: loaded.append(path))
    _stub_fsdp_loading(monkeypatch)
    return loaded


class _PackedTinyModel(torch.nn.Module):
    param_names_mapping = {r"^hf\.(.*)$": r"custom.\1"}

    def __init__(self, config=None):
        super().__init__()


def _packed_init_params():
    from types import SimpleNamespace
    return {"config": SimpleNamespace(quant_config=SimpleNamespace(layer_profile="h3_dit"))}


def test_packed_nvfp4_export_loads_with_a_low_rank_only_adapter(monkeypatch, tmp_path):
    loaded = _stub_packed_nvfp4_export(monkeypatch, tmp_path)
    monkeypatch.setattr(DenseLoRAPatch, "from_adapter", classmethod(lambda *_args, **_kwargs: None))

    _load_tiny_model(_PackedTinyModel, lora_path="adapter.safetensors", init_params=_packed_init_params())

    assert loaded == [str(tmp_path / "packed.safetensors")]


def test_packed_nvfp4_export_rejects_an_adapter_with_dense_payload(monkeypatch, tmp_path):
    loaded = _stub_packed_nvfp4_export(monkeypatch, tmp_path)
    patch = DenseLoRAPatch([], {"linear.weight": ("adapter.safetensors", "linear.diff")}, {}, 1.0)
    monkeypatch.setattr(DenseLoRAPatch, "from_adapter", classmethod(lambda *_args, **_kwargs: patch))

    with pytest.raises(ValueError, match="dense"):
        _load_tiny_model(_PackedTinyModel, lora_path="adapter.safetensors", init_params=_packed_init_params())

    assert loaded == []