"""Convert pinned source weights to strict MLX FP32 bundles. No training or quantization."""
import argparse
import importlib.metadata
import json
import platform
import shutil
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "scripts"))
import numpy as np
import mlx.core as mx
from safetensors.numpy import load_file
from timestral_mlx.config import TimestralConfig
from timestral_mlx.model import TimestralModel
from download import VARIANTS, sha256


def convert(variant):
    src, dst = ROOT / "models" / "source" / variant, ROOT / "models" / variant
    source = json.loads((src / "source.json").read_text())
    for name, entry in source["files"].items():
        if sha256(src / name) != entry["sha256"]:
            raise ValueError(f"source integrity check failed: {name}")
    dst.mkdir(parents=True, exist_ok=True)
    weights = load_file(src / "model.safetensors")
    if any(v.dtype != np.float32 for v in weights.values()):
        raise ValueError("expected original FP32 weights")
    converted = {k.replace("._gate.", ".gate."): mx.array(v) for k, v in weights.items()}
    config = TimestralConfig.load(src / "config.json")
    model = TimestralModel(config)
    model.load_weights(list(converted.items()), strict=True)
    mx.eval(model.parameters())
    mx.save_safetensors(str(dst / "model.safetensors"), converted,
                        metadata={"format": "mlx", "source": source["repo_id"], "revision": source["revision"]})
    # Check every serialized tensor, including exact source values and key mapping.
    restored = mx.load(str(dst / "model.safetensors"))
    if set(restored) != set(converted):
        raise ValueError("serialized key inventory mismatch")
    for k, v in weights.items():
        if not np.array_equal(v, np.array(restored[k.replace("._gate.", ".gate.")])):
            raise ValueError(f"lossless conversion check failed: {k}")
    shutil.copyfile(src / "config.json", dst / "config.json")
    shutil.copyfile(src / "LICENSE", dst / "LICENSE")
    manifest = {
        "format": "timestral-mlx-fp32-v1", "variant": variant, "dtype": "float32",
        "source": source, "source_code": {"url": "https://github.com/mistralai/timestral",
        "revision": "a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff"},
        "conversion": {"tensor_count": len(weights), "parameter_count": sum(v.size for v in weights.values()),
        "mapping": "Replace ._gate. with .gate.; tensor layouts and FP32 values unchanged",
        "lossless_roundtrip_verified": True, "training": False, "quantization": False},
        "environment": {"python": platform.python_version(), "mlx": importlib.metadata.version("mlx"),
                         "numpy": np.__version__, "platform": platform.platform()},
        "files": {name: {"sha256": sha256(dst / name), "bytes": (dst / name).stat().st_size}
                  for name in ("config.json", "model.safetensors", "LICENSE")},
    }
    (dst / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(f"Converted {variant}: {len(weights)} tensors, {manifest['conversion']['parameter_count']:,} parameters; lossless FP32 verified", flush=True)


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("variants", nargs="*", choices=VARIANTS, default=list(VARIANTS))
    for size in p.parse_args().variants:
        convert(size)