#!/usr/bin/env python3
# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Convert an official Modilify Mk1 checkpoint into native MLX shards."""

from __future__ import annotations

import argparse
import json
import shutil
import sys
from pathlib import Path

import mlx.core as mx

ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from modilify_mlx.config import MODEL_TYPE, TEXT_MODEL_TYPE, ModilifyMk1Config
from modilify_mlx.convert_utils import remap_weight, should_keep_source_key

COPY_PATTERNS = (
    "tokenizer.json",
    "tokenizer_config.json",
    "chat_template.jinja",
    "processor_config.json",
    "generation_config.json",
    "special_tokens_map.json",
    "preprocessor_config.json",
)

MAX_SHARD_BYTES = 5 * 1024**3


def _load_json(path: Path) -> dict:
    with path.open(encoding="utf-8") as handle:
        return json.load(handle)


def _write_config(source: Path, destination: Path) -> None:
    payload = _load_json(source / "config.json")
    if payload.get("model_type") != MODEL_TYPE:
        raise ValueError(
            f"Source model_type must be {MODEL_TYPE!r}, got {payload.get('model_type')!r}."
        )
    payload.pop("auto_map", None)
    payload["architectures"] = ["ModilifyMk1ForBlockDiffusion"]
    text = dict(payload.get("text_config") or {})
    text["model_type"] = TEXT_MODEL_TYPE
    payload["text_config"] = text
    payload["model_type"] = MODEL_TYPE
    # Validate through the native config so a bad export fails here.
    ModilifyMk1Config.from_dict(payload)
    with (destination / "config.json").open("w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2)
        handle.write("\n")


def _copy_sidecar_files(source: Path, destination: Path) -> None:
    for name in COPY_PATTERNS:
        src = source / name
        if src.exists():
            shutil.copy2(src, destination / name)


def _flush_shard(
    destination: Path,
    shard: dict[str, mx.array],
    shard_index: int,
    planned_count: int,
    weight_map: dict[str, str],
) -> int:
    if not shard:
        return shard_index
    name = f"model-{shard_index:05d}-of-{planned_count:05d}.safetensors"
    mx.save_safetensors(
        str(destination / name),
        shard,
        metadata={"format": "mlx"},
    )
    for key in shard:
        weight_map[key] = name
    return shard_index + 1


def convert(source: Path, destination: Path) -> None:
    source = source.resolve()
    destination = destination.resolve()
    destination.mkdir(parents=True, exist_ok=True)
    index = _load_json(source / "model.safetensors.index.json")
    source_map: dict[str, str] = index["weight_map"]
    shard_names = []
    for name in source_map.values():
        if name not in shard_names:
            shard_names.append(name)

    remapped: dict[str, mx.array] = {}
    kept_source = 0
    dropped_source = 0
    for shard_name in shard_names:
        print(f"[convert] reading {shard_name}", flush=True)
        loaded = mx.load(str(source / shard_name))
        for key, value in loaded.items():
            if not should_keep_source_key(key):
                dropped_source += 1
                continue
            kept_source += 1
            for new_key, new_value in remap_weight(key, value):
                if new_key in remapped:
                    raise ValueError(f"Duplicate remapped key: {new_key}")
                remapped[new_key] = new_value
        del loaded

    print(
        f"[convert] kept {kept_source} source tensors, dropped {dropped_source}, "
        f"wrote {len(remapped)} MLX tensors",
        flush=True,
    )
    latent_keys = [key for key in remapped if key.startswith("latent_deliberation.")]
    if len(latent_keys) < 150:
        raise RuntimeError(
            f"Expected the latent stack to survive conversion, found {len(latent_keys)} keys."
        )

    items = sorted(remapped.items())
    shards: list[dict[str, mx.array]] = []
    current: dict[str, mx.array] = {}
    current_bytes = 0
    for key, value in items:
        tensor_bytes = int(value.nbytes)
        if current and current_bytes + tensor_bytes > MAX_SHARD_BYTES:
            shards.append(current)
            current = {}
            current_bytes = 0
        current[key] = value
        current_bytes += tensor_bytes
    if current:
        shards.append(current)

    planned = max(1, len(shards))
    weight_map: dict[str, str] = {}
    total_size = 0
    for index_i, shard in enumerate(shards, start=1):
        print(f"[convert] writing shard {index_i}/{planned}", flush=True)
        name = (
            f"model-{index_i:05d}-of-{planned:05d}.safetensors"
            if planned > 1
            else "model.safetensors"
        )
        mx.save_safetensors(
            str(destination / name),
            shard,
            metadata={"format": "mlx"},
        )
        for key, value in shard.items():
            weight_map[key] = name
            total_size += int(value.nbytes)

    index_payload = {
        "metadata": {"total_size": total_size},
        "weight_map": {key: weight_map[key] for key in sorted(weight_map)},
    }
    with (destination / "model.safetensors.index.json").open(
        "w", encoding="utf-8"
    ) as handle:
        json.dump(index_payload, handle, indent=2)
        handle.write("\n")

    _write_config(source, destination)
    _copy_sidecar_files(source, destination)
    print(f"[convert] done -> {destination}", flush=True)
    print(f"[convert] tensors={len(weight_map)} bytes={total_size}", flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--source",
        type=Path,
        default=Path.home() / "Modilify-Mk1",
        help="Official Mk1 Hugging Face directory",
    )
    parser.add_argument(
        "--destination",
        type=Path,
        default=ROOT,
        help="Native MLX output directory",
    )
    args = parser.parse_args()
    convert(args.source, args.destination)


if __name__ == "__main__":
    main()