#
#  zyquo_fuse.py
#  Zyquo MLX
#
#  Author: Simon-Pierre Boucher
#  Mail: contact@spboucher.ai
#
#  Adapter-fuse driver: fuses a LoRA adapter into its base model and saves a
#  standalone MLX model directory (mlx_lm.fuse semantics —
#  docs/TRAINING-RESEARCH.md §2.2). JSON-lines protocol on stdout.
#
#  Usage: zyquo_fuse.py --model <base> --adapter-path <dir> --save-path <dir>
#         [--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("--model", required=True)
    parser.add_argument("--adapter-path", required=True)
    parser.add_argument("--save-path", required=True)
    parser.add_argument("--dequantize", action="store_true")
    args = parser.parse_args()

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

        # mlx_lm.fuse exposes only a CLI main(); call it with argv to stay
        # on the supported surface of the pinned version.
        from mlx_lm import fuse

        argv = [
            "--model", args.model,
            "--adapter-path", args.adapter_path,
            "--save-path", args.save_path,
        ]
        if args.dequantize:
            argv.append("--dequantize")  # exact flag name in mlx-lm 0.31.3

        old_argv = sys.argv
        sys.argv = ["mlx_lm.fuse"] + argv
        try:
            fuse.main()
        finally:
            sys.argv = old_argv

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


if __name__ == "__main__":
    main()