import json

import mlx.core as mx
import numpy as np
import pytest
import torch
from safetensors.torch import save_file

from h3_mlx.convert import convert
from h3_mlx.weights import Linear, WeightStore


def test_quantized_shards_roundtrip_and_preserve_qkv_layout(tmp_path):
    model = tmp_path / "original"
    transformer = model / "transformer"
    transformer.mkdir(parents=True)
    config = {
        "hidden_size": 128,
        "num_layers": 1,
        "token_refiner_num_layers": 0,
        "num_attention_heads": 2,
        "attention_head_dim": 128,
        "ffn_hidden_size": 128,
        "audio_enabled": False,
    }
    (transformer / "config.json").write_text(json.dumps(config))
    torch.manual_seed(77)
    tensors = {}
    for name in (
        "norm1.weight",
        "norm2.weight",
        "attn.q_norm.weight",
        "attn.k_norm.weight",
    ):
        tensors["blocks.0." + name] = torch.randn(128).bfloat16()
    for name, shape in (
        ("attn.qkv_proj", (768, 128)),
        ("attn.out_proj", (128, 256)),
        ("mlp.fc1", (256, 128)),
        ("mlp.fc2", (128, 128)),
    ):
        tensors[f"blocks.0.{name}.weight"] = torch.randn(*shape).bfloat16()
    save_file(tensors, str(transformer / "model.safetensors"))
    output = tmp_path / "converted"
    meta = convert(model, output, bits=8, group_size=64)
    assert meta["format"] == "h3-mlx-core-v1" and len(meta["files"]) == 1
    original = WeightStore(transformer)
    converted = WeightStore(output)
    a = Linear.load(original, "blocks.0.attn.qkv_proj", bits=8, grouped_heads=(2, 128))
    b = Linear.load(converted, "blocks.0.attn.qkv_proj", grouped_heads=(2, 128))
    for x, y in ((a.weight, b.weight), (a.scales, b.scales), (a.qbiases, b.qbiases)):
        np.testing.assert_array_equal(
            np.asarray(x.view(mx.uint16) if x.dtype == mx.bfloat16 else x),
            np.asarray(y.view(mx.uint16) if y.dtype == mx.bfloat16 else y),
        )
    with pytest.raises(FileExistsError):
        convert(model, output)