"""Save and reload MLX packs with an explicit matching-runtime requirement."""

import json
import shutil
from pathlib import Path

import mlx.core as mx
from mlx.utils import tree_flatten
from mlx_lm.models.qwen3_5 import TextModel, TextModelArgs
from runtime import Packed

from mlx import nn


def save_model(model, info, tokenizer_path, directory):
    directory = Path(directory)
    modules = []
    for name, module in model.named_modules():
        if isinstance(module, Packed):
            if module.dtype != mx.float16:
                raise ValueError("Only float16 packed artifacts are supported")
            modules.append(
                {
                    "path": name,
                    "block": module.block,
                    "embedding": module.embedding,
                    "dtype": "float16",
                }
            )
    config = {
        "schema_version": 1,
        "model_type": "prism_hadamard_qwen35",
        "requires_runtime": "examples/python/hadamard/artifact.py",
        "text_config": info["config"],
        "modules": modules,
        "quantization": {"bits": 2, "group_size": 128, "mode": "affine"},
    }
    directory.mkdir(parents=True, exist_ok=False)
    mx.save_safetensors(
        str(directory / "model.safetensors"),
        dict(tree_flatten(model.parameters())),
        metadata={"format": "mlx"},
    )
    (directory / "config.json").write_text(json.dumps(config, indent=2))
    shutil.copyfile(tokenizer_path, directory / "tokenizer.json")
    (directory / "README.md").write_text(
        "# Experimental folded MLX pack\n\nRequires the accompanying Hadamard-aware loader. Do not load as an ordinary affine Qwen checkpoint: that would omit activation and inverse-embedding transforms. The source numerical and performance reports are separate from serialization integrity.\n"
    )
    return config


def validate_record(original, record, arrays, signs):
    if not isinstance(original, (nn.Linear, nn.Embedding)):
        raise ValueError("Unsupported packed module target")
    if record["embedding"] != isinstance(original, nn.Embedding):
        raise ValueError("Packed module kind mismatch")
    rows, width = original.weight.shape
    if width % 128:
        raise ValueError("Invalid packed width")
    expected = [(rows, width // 16), (rows, width // 128), (rows, width // 128)]
    if [a.shape for a in arrays] != expected or arrays[0].dtype != mx.uint32:
        raise ValueError("Invalid packed tensor shapes or storage dtype")
    for array in arrays[1:]:
        if array.dtype not in (mx.float16, mx.float32, mx.bfloat16):
            raise ValueError("Invalid affine dtype")
        if not mx.all(mx.isfinite(array)).item():
            raise ValueError("Non-finite affine parameters")
    block = record["block"]
    if block:
        if width % block or signs is None or signs.shape != (width,):
            raise ValueError("Invalid transform dimensions")
        if not mx.all((signs == 1) | (signs == -1)).item():
            raise ValueError("Invalid sign values")
    elif signs is not None:
        raise ValueError("Unexpected sign vector")


def load_model(directory):
    directory = Path(directory)
    config = json.loads((directory / "config.json").read_text())
    if (
        config.get("schema_version") != 1
        or config.get("model_type") != "prism_hadamard_qwen35"
    ):
        raise ValueError("Unsupported packed model schema")
    model = TextModel(TextModelArgs.from_dict(config["text_config"]))
    weights = mx.load(str(directory / "model.safetensors"))
    seen = set()
    for record in config["modules"]:
        if record["path"] in seen:
            raise ValueError("Duplicate packed module")
        seen.add(record["path"])
        parts = record["path"].split(".")
        parent = model
        for part in parts[:-1]:
            parent = parent[int(part)] if part.isdigit() else getattr(parent, part)
        name = record["path"]
        arrays = [
            weights[name + "." + suffix] for suffix in ("weight", "scales", "biases")
        ]
        if record["dtype"] != "float16":
            raise ValueError("Unsupported activation dtype")
        block = record["block"]
        if block and block not in (512, 1024, 2048, 4096):
            raise ValueError("Unsupported block size")
        signs = weights.get(name + ".signs")
        if block and signs is None:
            raise ValueError("Missing sign vector")
        original = getattr(parent, parts[-1])
        validate_record(original, record, arrays, signs)
        setattr(
            parent,
            parts[-1],
            Packed(arrays, block, signs, record["embedding"], mx.float16),
        )
    model.load_weights(list(weights.items()), strict=True)
    model.eval()
    mx.eval(model.parameters())
    return model, config


def save_and_verify(model, info, tokenizer_path, directory, input_ids):
    import numpy as np

    config = save_model(model, info, tokenizer_path, directory)
    reloaded, _ = load_model(directory)
    x = mx.array([input_ids], dtype=mx.int32)

    def evaluate(m):
        hidden = m.model(x, cache=m.make_cache())
        y = m.lm_head(hidden[:, -1:, :])
        mx.eval(y)
        return np.asarray(y)

    expected = evaluate(model)
    actual = evaluate(reloaded)
    np.testing.assert_array_equal(actual, expected)
    result = {
        "reload_logits_exact": True,
        "checked_logits": actual.size,
        "packed_modules": len(config["modules"]),
    }
    Path(directory, "reload-validation.json").write_text(json.dumps(result, indent=2))
    return result