#!/usr/bin/env python3
"""
Split a single MLX model.safetensors into sharded safetensors for
HuggingFace upload, generating model.safetensors.index.json.

Preserves dtypes (bf16, uint32 quant packs, etc.) since it loads and
saves natively with MLX.
"""

import argparse
import json
import shutil
from pathlib import Path

import mlx.core as mx


def human_to_bytes(s: str) -> int:
    s = s.strip().upper()
    units = {"B": 1, "KB": 1024, "MB": 1024**2,
             "GB": 1024**3, "TB": 1024**4}
    for u in ("TB", "GB", "MB", "KB", "B"):
        if s.endswith(u):
            return int(float(s[:-len(u)]) * units[u])
    return int(s)  # raw byte count


def dtype_size(arr: mx.array) -> int:
    """Bytes per element for the array's dtype."""
    # mx.array.nbytes gives total bytes directly.
    return arr.nbytes


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", "-i", required=True,
                    help="Merged model dir OR path to a single "
                         "model.safetensors")
    ap.add_argument("--out", "-o", required=True,
                    help="Output directory for sharded model")
    ap.add_argument("--max-shard-size", default="5GB",
                    help="Max size per shard (e.g. 5GB, 4GB). Default 5GB.")
    ap.add_argument("--copy-aux", action="store_true",
                    help="Copy config/tokenizer/etc. from input dir to out.")
    args = ap.parse_args()

    in_path = Path(args.input)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    # Resolve the single-file source and its parent dir (for aux files).
    if in_path.is_dir():
        src_file = in_path / "model.safetensors"
        src_dir = in_path
        if not src_file.exists():
            # Maybe it's already sharded — bail with guidance.
            shards = sorted(in_path.glob("model-*.safetensors"))
            if shards:
                print("Input dir already contains sharded safetensors:")
                for s in shards:
                    print("  ", s.name)
                print("This script expects a single model.safetensors. "
                      "Point --input at that file, or consolidate first.")
                return
            raise FileNotFoundError(f"No model.safetensors in {in_path}")
    else:
        src_file = in_path
        src_dir = in_path.parent

    max_bytes = human_to_bytes(args.max_shard_size)
    print(f"Loading {src_file} …")
    weights = mx.load(str(src_file))
    print(f"Loaded {len(weights)} tensors.")

    # Compute total size and per-tensor sizes.
    sizes = {k: dtype_size(v) for k, v in weights.items()}
    total = sum(sizes.values())
    print(f"Total weight size: {total / 1024**3:.2f} GB")
    print(f"Target max shard size: {max_bytes / 1024**3:.2f} GB")

    # Greedy bin-packing into shards, preserving insertion order.
    # (Keeps related tensors together reasonably well.)
    shards = []           # list of dict[name -> array]
    current = {}
    current_size = 0

    for name, arr in weights.items():
        sz = sizes[name]
        if sz > max_bytes:
            # A single tensor exceeds the shard limit; it gets its own shard.
            if current:
                shards.append(current)
                current, current_size = {}, 0
            shards.append({name: arr})
            print(f"  NOTE: '{name}' ({sz/1024**3:.2f} GB) exceeds shard "
                  f"limit; placed in its own shard.")
            continue
        if current_size + sz > max_bytes and current:
            shards.append(current)
            current, current_size = {}, 0
        current[name] = arr
        current_size += sz
    if current:
        shards.append(current)

    n = len(shards)
    print(f"Splitting into {n} shard(s).")

    if n == 1:
        # Single shard: HF convention is just model.safetensors (no index).
        out_file = out / "model.safetensors"
        mx.save_safetensors(str(out_file), shards[0],
                            metadata={"format": "mlx"})
        print(f"Wrote {out_file.name} (single shard, no index needed).")
    else:
        # Multi-shard: model-00001-of-000NN.safetensors + index.
        weight_map = {}
        for i, shard in enumerate(shards, start=1):
            fname = f"model-{i:05d}-of-{n:05d}.safetensors"
            mx.save_safetensors(str(out / fname), shard,
                                metadata={"format": "mlx"})
            for k in shard:
                weight_map[k] = fname
            shard_bytes = sum(sizes[k] for k in shard)
            print(f"  Wrote {fname} "
                  f"({len(shard)} tensors, {shard_bytes/1024**3:.2f} GB)")

        index = {
            "metadata": {"total_size": total},
            "weight_map": weight_map,
        }
        idx_file = out / "model.safetensors.index.json"
        idx_file.write_text(json.dumps(index, indent=2))
        print(f"Wrote {idx_file.name}")

    if args.copy_aux:
        copy_aux(src_dir, out)

    print("\nDone. Upload with:")
    print(f"  huggingface-cli upload <repo_id> {out} .")


def copy_aux(src_dir: Path, out: Path):
    aux = [
        "config.json",
        "tokenizer.json", "tokenizer_config.json", "vocab.json",
        "merges.txt", "special_tokens_map.json", "added_tokens.json",
        "chat_template.jinja", "generation_config.json",
        "preprocessor_config.json", "processor_config.json",
        "image_processor_config.json", "video_processor_config.json",
    ]
    copied = 0
    for n in aux:
        src = src_dir / n
        if src.exists():
            shutil.copy(src, out / n)
            copied += 1
    print(f"aux: copied {copied} auxiliary file(s) from {src_dir}")


if __name__ == "__main__":
    main()