"""Keep serving LoRA merge consistent with exported training parameters."""

import json

import pytest
import torch

from wall_x._vendor.harrix.utils import ckpt_load


def _tiny_lora():
    stem = "model.base_model.model.proj"
    return {
        stem + ".base_layer.weight": torch.zeros((2, 2)),
        stem + ".lora_A.default.weight": torch.eye(2),
        stem + ".lora_B.default.weight": torch.eye(2),
    }


def test_checkpoint_local_lora_config_controls_merge(tmp_path):
    (tmp_path / "lora_config.json").write_text(
        json.dumps({"lora_r": 2, "lora_alpha": 6}), encoding="utf-8"
    )
    tensors = _tiny_lora()
    scale = ckpt_load.resolve_lora_scale(str(tmp_path), {}, tensors)
    merged = ckpt_load.reshape_compatible_state_dict(
        tensors, {"model.proj.weight": torch.zeros((2, 2))}, lora_scale=scale
    )
    assert scale == 3
    torch.testing.assert_close(merged["model.proj.weight"], 3 * torch.eye(2))


def test_checkpoint_lora_rank_mismatch_is_rejected(tmp_path):
    (tmp_path / "lora_config.json").write_text(
        json.dumps({"lora_r": 4, "lora_alpha": 8}), encoding="utf-8"
    )
    with pytest.raises(ValueError, match="rank"):
        ckpt_load.resolve_lora_scale(str(tmp_path), {}, _tiny_lora())


def test_legacy_checkpoint_without_lora_metadata_uses_previous_scale(tmp_path):
    messages = []
    scale = ckpt_load.resolve_lora_scale(str(tmp_path), {}, _tiny_lora(), log_fn=messages.append)
    assert scale == 2
    assert any("metadata" in message for message in messages)