import json

import mlx.core as mx
import pytest

from ark_asr_mlx.conversion import load_source_weights, sanitize_weights


def test_sanitize_transposes_conv_and_remaps_adapter() -> None:
    source = {
        "audio_encoder.whisper.conv1.weight": mx.zeros((8, 4, 3)),
        "audio_encoder.adapting.0.weight": mx.zeros((16, 32)),
        "lm_head.weight": mx.zeros((20, 8)),
        "audio_encoder.whisper.embed_positions.weight": mx.zeros((10, 8)),
    }
    converted, dropped = sanitize_weights(
        source,
        tied_embeddings=True,
        use_audio_rope=True,
    )
    assert converted["audio_encoder.whisper.conv1.weight"].shape == (8, 3, 4)
    assert "audio_encoder.linear1.weight" in converted
    assert "lm_head.weight" in dropped
    assert "audio_encoder.whisper.embed_positions.weight" in dropped


def test_load_source_weights_from_single_file(tmp_path) -> None:
    mx.save_safetensors(
        str(tmp_path / "model.safetensors"),
        {"model.embed_tokens.weight": mx.zeros((4, 2))},
    )

    weights, weight_files = load_source_weights(tmp_path)

    assert set(weights) == {"model.embed_tokens.weight"}
    assert [path.name for path in weight_files] == ["model.safetensors"]


def test_load_source_weights_from_indexed_shards(tmp_path) -> None:
    first_name = "model-00001-of-00002.safetensors"
    second_name = "model-00002-of-00002.safetensors"
    mx.save_safetensors(str(tmp_path / first_name), {"first.weight": mx.zeros((2, 2))})
    mx.save_safetensors(str(tmp_path / second_name), {"second.weight": mx.zeros((3, 2))})
    (tmp_path / "model.safetensors.index.json").write_text(
        json.dumps(
            {
                "metadata": {"total_size": 40},
                "weight_map": {
                    "first.weight": first_name,
                    "second.weight": second_name,
                },
            }
        ),
        encoding="utf-8",
    )

    weights, weight_files = load_source_weights(tmp_path)

    assert set(weights) == {"first.weight", "second.weight"}
    assert [path.name for path in weight_files] == [first_name, second_name]


def test_load_source_weights_rejects_missing_shard(tmp_path) -> None:
    (tmp_path / "model.safetensors.index.json").write_text(
        json.dumps({"weight_map": {"missing.weight": "missing.safetensors"}}),
        encoding="utf-8",
    )

    with pytest.raises(FileNotFoundError, match="missing.safetensors"):
        load_source_weights(tmp_path)