"""Stream the official checkpoint into MLX-VLM-compatible 4/8-bit shards."""
import argparse
import copy
import json
import shutil
import struct
import time
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import numpy as np
from mlx.utils import tree_flatten
from mlx_vlm.models.muse_glimmer import Model, ModelConfig
from mlx_vlm.utils import skip_multimodal_module


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--revision", required=True)
    parser.add_argument("--wait-for-download", action="store_true")
    args = parser.parse_args()
    config = json.loads((args.source / "config.json").read_text())
    model = Model(ModelConfig.from_dict(config))
    shapes = {k: v.shape for k, v in tree_flatten(model.parameters())}
    quantized_paths = {
        path for path, module in model.named_modules()
        if isinstance(module, (nn.Linear, nn.Embedding))
        and not skip_multimodal_module(path)
        and module.weight.shape[-1] % 64 == 0
    }
    source_index = json.loads((args.source / "model.safetensors.index.json").read_text())
    mapped = model.sanitize({key: None for key in source_index["weight_map"]})
    if set(mapped) != set(shapes):
        raise ValueError(f"Schema mismatch: missing={set(shapes)-set(mapped)}, extra={set(mapped)-set(shapes)}")
    del model
    mx.clear_cache()
    mx.set_cache_limit(256 * 1024**2)
    states = {}
    for bits in (4, 8):
        folder = args.output / f"meta-encoder-mlx-{bits}bit"
        folder.mkdir(parents=True, exist_ok=True)
        if list(folder.glob("*.safetensors")):
            raise FileExistsError(f"Refusing to overwrite existing weights in {folder}")
        states[bits] = dict(folder=folder, pending={}, bytes=0, total=0, index={}, shards=[])

    def flush(state):
        if not state["pending"]:
            return
        name = f"part-{len(state['shards'])+1:05d}.safetensors"
        mx.save_safetensors(str(state["folder"] / name), state["pending"], metadata={"format": "mlx"})
        state["index"].update({key: name for key in state["pending"]})
        state["shards"].append(name)
        state["pending"].clear()
        state["bytes"] = 0

    count = 0
    for shard_name in sorted(set(source_index["weight_map"].values()), reverse=True):
        if args.wait_for_download:
            while not (args.source / shard_name).exists():
                print(f"Waiting for download: {shard_name}", flush=True)
                time.sleep(30)
        with (args.source / shard_name).open("rb") as source:
            header_size = struct.unpack("<Q", source.read(8))[0]
            header = json.loads(source.read(header_size))
            for original, info in header.items():
                if original == "__metadata__":
                    continue
                key = next(iter(Model.sanitize(None, {original: None})))
                if tuple(info["shape"]) != tuple(shapes[key]):
                    raise ValueError(f"Shape mismatch for {key}")
                if info["dtype"] != "BF16":
                    raise ValueError(f"Unexpected dtype {info['dtype']} for {key}")
                start, end = info["data_offsets"]
                source.seek(8 + header_size + start)
                raw = source.read(end - start)
                if len(raw) != end - start:
                    raise EOFError(original)
                weight = mx.array(np.frombuffer(raw, dtype=np.uint16).reshape(info["shape"])).view(mx.bfloat16)
                mx.eval(weight)
                del raw
                for bits, state in states.items():
                    if key.removesuffix(".weight") in quantized_paths and key.endswith(".weight"):
                        packed, scales, biases = mx.quantize(weight, group_size=64, bits=bits, mode="affine")
                        arrays = {key: packed, key.removesuffix(".weight") + ".scales": scales,
                                  key.removesuffix(".weight") + ".biases": biases}
                    else:
                        arrays = {key: weight}
                    mx.eval(list(arrays.values()))
                    size = sum(value.nbytes for value in arrays.values())
                    if state["bytes"] + size > 2 * 1024**3:
                        flush(state)
                    state["pending"].update(arrays)
                    state["bytes"] += size
                    state["total"] += size
                del arrays, weight
                count += 1
                if count % 100 == 0:
                    print(f"Converted {count}/{len(shapes)} tensors", flush=True)
    for bits, state in states.items():
        flush(state)
        names = {}
        for number, old in enumerate(state["shards"], 1):
            new = f"model-{number:05d}-of-{len(state['shards']):05d}.safetensors"
            (state["folder"] / old).rename(state["folder"] / new)
            names[old] = new
        index = {"metadata": {"total_size": state["total"]},
                 "weight_map": {k: names[v] for k, v in sorted(state["index"].items())}}
        (state["folder"] / "model.safetensors.index.json").write_text(json.dumps(index, indent=2) + "\n")
        for name in ("tokenizer.json", "tokenizer_config.json", "processor_config.json",
                     "generation_config.json", "chat_template.jinja", "LICENSE", "NOTICE", "USAGE_POLICY.md"):
            if (args.source / name).exists():
                shutil.copy2(args.source / name, state["folder"] / name)
        out_config = copy.deepcopy(config)
        out_config["quantization"] = {"group_size": 64, "bits": bits, "mode": "affine"}
        out_config["quantization_config"] = copy.deepcopy(out_config["quantization"])
        out_config["vision_config"]["skip_vision"] = True
        (state["folder"] / "config.json").write_text(json.dumps(out_config, indent=2) + "\n")
        provenance = {"source": "facebook/meta-encoder", "revision": args.revision,
                      "bits": bits, "group_size": 64, "mode": "affine",
                      "vision": "unchanged BF16", "source_tensors": count,
                      "quantized_modules": len(quantized_paths), "weight_bytes": state["total"],
                      "mlx": mx.__version__}
        (state["folder"] / "conversion.json").write_text(json.dumps(provenance, indent=2) + "\n")
        print(json.dumps(provenance), flush=True)


if __name__ == "__main__":
    main()