#
#  zyquo_convert.py
#  Zyquo MLX
#
#  Author: Simon-Pierre Boucher
#  Mail: contact@spboucher.ai
#
#  Conversion/quantization driver: HF model → MLX format, optionally
#  quantized (mlx_lm.convert semantics — docs/MLX-RESEARCH.md §5.2).
#  JSON-lines protocol on stdout.
#
#  Usage: zyquo_convert.py --hf-path <repo-or-dir> --mlx-path <dir>
#         [--quantize --q-bits 4 --q-group-size 64 --q-mode affine]
#         [--dtype float16|bfloat16|float32] [--dequantize]

import argparse
import json
import sys
import time


def emit(obj):
    sys.stdout.write(json.dumps(obj) + "\n")
    sys.stdout.flush()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--hf-path", required=True)
    parser.add_argument("--mlx-path", required=True)
    parser.add_argument("--quantize", action="store_true")
    parser.add_argument("--q-bits", type=int, default=None)
    parser.add_argument("--q-group-size", type=int, default=None)
    parser.add_argument("--q-mode", default="affine",
                        choices=["affine", "mxfp4", "nvfp4", "mxfp8"])
    parser.add_argument("--dtype", default=None,
                        choices=["float16", "bfloat16", "float32"])
    parser.add_argument("--dequantize", action="store_true")
    args = parser.parse_args()

    try:
        emit({"event": "start", "stage": "convert", "hf_path": args.hf_path, "ts": time.time()})

        from mlx_lm import convert

        # mlx-lm 0.31.3 first downloads weights with allow_patterns, then its
        # copy step demands the FULL snapshot with local_files_only=True and
        # trips IncompleteSnapshotError on aux files (.gitattributes, LICENSE…).
        # Pre-fetch the complete snapshot for remote repos to sidestep it.
        import os
        if not os.path.isdir(args.hf_path):
            from huggingface_hub import snapshot_download
            emit({"event": "stage", "stage": "downloading"})
            snapshot_download(repo_id=args.hf_path)

        kwargs = {
            "hf_path": args.hf_path,
            "mlx_path": args.mlx_path,
            "quantize": args.quantize,
            "dequantize": args.dequantize,
        }
        if args.q_bits is not None:
            kwargs["q_bits"] = args.q_bits
        if args.q_group_size is not None:
            kwargs["q_group_size"] = args.q_group_size
        if args.quantize:
            kwargs["q_mode"] = args.q_mode
        if args.dtype:
            kwargs["dtype"] = args.dtype

        convert(**kwargs)
        emit({"event": "done", "mlx_path": args.mlx_path})
    except Exception as exc:  # noqa: BLE001
        emit({"event": "error", "message": str(exc), "type": type(exc).__name__})
        sys.exit(1)


if __name__ == "__main__":
    main()