# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
from unittest.mock import Mock

import pytest
import torch
import torch.nn as nn
from safetensors.torch import save_file
from transformers import CLIPTextConfig, CLIPTextModel

from diffusers.configuration_utils import ConfigMixin
from diffusers.loaders import StableDiffusionLoraLoaderMixin, lora_base
from diffusers.loaders.lora_base import LoraBaseMixin
from diffusers.loaders.peft import PeftAdapterMixin
from diffusers.models.modeling_utils import ModelMixin
from diffusers.utils import is_transformers_version
from diffusers.utils.import_utils import is_peft_available

from ..testing_utils import CaptureLogger, require_peft_backend


if is_peft_available():
    from peft import LoraConfig
    from peft.tuners.tuners_utils import BaseTunerLayer


LORA_KEY = "unet.test.lora_A.weight"


def _write_lora_weights(path):
    save_file({LORA_KEY: torch.ones(1)}, path)


@pytest.fixture
def lora_weight_path(tmp_path):
    weight_path = tmp_path / "adapter.safetensors"
    _write_lora_weights(weight_path)
    return weight_path


@pytest.fixture
def model_info_mock(monkeypatch):
    model_info_mock = Mock()
    monkeypatch.setattr(lora_base, "model_info", model_info_mock)
    return model_info_mock


def test_local_directory_in_offline_mode(lora_weight_path, monkeypatch, model_info_mock):
    monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)

    state_dict, _ = StableDiffusionLoraLoaderMixin.lora_state_dict(lora_weight_path.parent)

    assert torch.equal(state_dict[LORA_KEY], torch.ones(1))
    model_info_mock.assert_not_called()


def test_local_directory_with_local_files_only(lora_weight_path, model_info_mock):
    state_dict, _ = StableDiffusionLoraLoaderMixin.lora_state_dict(lora_weight_path.parent, local_files_only=True)

    assert torch.equal(state_dict[LORA_KEY], torch.ones(1))
    model_info_mock.assert_not_called()


def test_local_file_in_offline_mode(lora_weight_path, monkeypatch, model_info_mock):
    monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)

    state_dict, _ = StableDiffusionLoraLoaderMixin.lora_state_dict(lora_weight_path)

    assert torch.equal(state_dict[LORA_KEY], torch.ones(1))
    model_info_mock.assert_not_called()


def test_remote_repository_in_offline_mode_requires_weight_name(monkeypatch, model_info_mock):
    monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)

    with pytest.raises(ValueError, match="offline mode.*weight_name"):
        StableDiffusionLoraLoaderMixin.lora_state_dict("organization/repository")

    model_info_mock.assert_not_called()


def test_local_directory_without_matching_files_returns_none(tmp_path, monkeypatch):
    (tmp_path / "notes.txt").touch()
    monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)

    weight_name = lora_base._best_guess_weight_name(tmp_path)

    assert weight_name is None


def test_local_directory_with_multiple_files_warns_and_uses_first(tmp_path, monkeypatch):
    first_path = tmp_path / "first.safetensors"
    second_path = tmp_path / "second.safetensors"
    first_path.touch()
    second_path.touch()
    monkeypatch.setattr(lora_base, "HF_HUB_OFFLINE", True)
    # `os.listdir` returns entries in arbitrary order; pin it so `first.safetensors` is the one picked.
    monkeypatch.setattr(lora_base.os, "listdir", lambda _: [first_path.name, second_path.name])

    # `caplog` cannot see this warning: `diffusers` sets `propagate = False` on its library root logger, so
    # records never reach the root handler pytest installs. `CaptureLogger` attaches to the logger directly.
    lora_base.logger.setLevel(logging.WARNING)
    with CaptureLogger(lora_base.logger) as cap_logger:
        weight_name = lora_base._best_guess_weight_name(tmp_path)

    assert weight_name == first_path.name
    assert "contains more than one weights file" in cap_logger.out


@require_peft_backend
def test_unfuse_lora_partial_components_keeps_merged_adapters_in_sync():
    """Regression test for gh-14214.

    Unfusing only a subset of components must keep _merged_adapters in sync
    with the adapters still physically fused in the remaining components.
    """

    class TinyModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
        config_name = "config.json"

        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(8, 8)

    class FakePipeline(LoraBaseMixin):
        _lora_loadable_modules = ["unet", "text_encoder"]

        def __init__(self, unet, text_encoder):
            self._merged_adapters = set()
            self.unet, self.text_encoder = unet, text_encoder

    unet = TinyModel()
    text_encoder = TinyModel()
    config = LoraConfig(r=4, lora_alpha=4, target_modules=["linear"], init_lora_weights=False)
    unet.add_adapter(config, adapter_name="adapter")
    text_encoder.add_adapter(config, adapter_name="adapter")

    pipe = FakePipeline(unet, text_encoder)
    pipe.fuse_lora(components=["unet", "text_encoder"], adapter_names=["adapter"])
    assert pipe.num_fused_loras == 1

    pipe.unfuse_lora(components=["text_encoder"])
    assert "adapter" in pipe.fused_loras, "adapter should remain tracked while unet is still fused"
    assert pipe.num_fused_loras == 1

    unet_still_merged = any(isinstance(m, BaseTunerLayer) and len(m.merged_adapters) > 0 for m in unet.modules())
    assert unet_still_merged, "unet should still be physically merged at the PEFT level"

    pipe.unfuse_lora(components=["unet"])
    assert pipe.num_fused_loras == 0


@require_peft_backend
@pytest.mark.skipif(
    not is_transformers_version(">=", "5.6.0"),
    reason="CLIPTextModel uses a flattened module namespace starting with transformers 5.6",
)
def test_load_lora_into_flattened_clip_text_encoder():
    config = CLIPTextConfig(
        vocab_size=100,
        hidden_size=32,
        intermediate_size=64,
        num_hidden_layers=1,
        num_attention_heads=4,
        max_position_embeddings=77,
        bos_token_id=0,
        eos_token_id=1,
        pad_token_id=2,
    )
    text_encoder = CLIPTextModel(config)

    lora_a = torch.randn(4, 32)
    lora_b = torch.randn(32, 4)

    state_dict = {
        "text_encoder.text_model.encoder.layers.0.self_attn.q_proj.lora_A.weight": lora_a,
        "text_encoder.text_model.encoder.layers.0.self_attn.q_proj.lora_B.weight": lora_b,
    }

    StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder(
        state_dict=state_dict,
        network_alphas=None,
        text_encoder=text_encoder,
        prefix="text_encoder",
        adapter_name="test",
        low_cpu_mem_usage=False,
    )

    q_proj = text_encoder.encoder.layers[0].self_attn.q_proj

    torch.testing.assert_close(q_proj.lora_A["test"].weight, lora_a)
    torch.testing.assert_close(q_proj.lora_B["test"].weight, lora_b)