#!/usr/bin/env python3
"""
Merge Qwen3.5-MoE vision_tower (dir B) into the finetuned
language-only quantized model (dir A).

Verified clean namespace split:
  language_model.*  -> LM (from A, keeps A's quantization map)
  vision_tower.*    -> vision (from B, unquantized bf16)

Config: LM + quantization from A (finetune), vision fields from B.
Chat template: from B by default (vision-aware).
"""

import argparse
import json
import shutil
from pathlib import Path

import mlx.core as mx


def load_all(model_dir: Path) -> dict:
    """Load every tensor from all safetensors shards, natively in MLX.

    mx.load reads safetensors directly and preserves bfloat16 (unlike a
    NumPy round-trip, which has no native bf16 dtype).
    """
    weights = {}
    files = sorted(model_dir.glob("*.safetensors"))
    if not files:
        raise FileNotFoundError(f"No safetensors in {model_dir}")
    for f in files:
        weights.update(mx.load(str(f)))
    return weights


def is_lm_key(k: str) -> bool:
    return k.startswith("language_model.")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--finetune", "-a", required=True,
                    help="Dir A: finetuned language-only quantized model")
    ap.add_argument("--base", "-b", required=True,
                    help="Dir B: base VLM with vision_tower")
    ap.add_argument("--out", "-o", required=True,
                    help="Output directory for merged model")
    ap.add_argument("--template-from", choices=["a", "b"], default="a",
                    help="Which model's chat template to use "
                         "(default b: vision-aware).")
    args = ap.parse_args()

    dir_a, dir_b, out = Path(args.finetune), Path(args.base), Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    print("Loading A (finetuned LM)…")
    a = load_all(dir_a)
    print("Loading B (base VLM)…")
    b = load_all(dir_b)

    lm_weights = {k: v for k, v in a.items() if is_lm_key(k)}
    vision_weights = {k: v for k, v in b.items() if not is_lm_key(k)}

    # Verify the vision split is what we expect (all vision_tower.*).
    non_vision_b = [k for k in vision_weights
                    if not k.startswith("vision_tower.")]
    if non_vision_b:
        print("WARNING: B has non-LM keys that are NOT vision_tower.*:")
        for k in non_vision_b[:20]:
            print("   ", k)
        print("   -> Decide whether these belong in the merge. "
              "Currently they WILL be included.")

    stray_a = [k for k in a if not is_lm_key(k)]
    if stray_a:
        print(f"NOTE: A has {len(stray_a)} non-LM keys (ignored). "
              f"e.g. {stray_a[:5]}")

    # Drop MTP tensors if present in B (A config sets mtp layers = 0).
    mtp = [k for k in vision_weights if "mtp" in k.lower()]
    for k in mtp:
        vision_weights.pop(k)
    if mtp:
        print(f"Dropped {len(mtp)} MTP tensors from B.")

    print(f"LM tensors (A):     {len(lm_weights)}")
    print(f"Vision tensors (B): {len(vision_weights)}")
    if not vision_weights:
        print("ERROR: no vision tensors. Aborting.")
        return

    merged = {**lm_weights, **vision_weights}
    print(f"Total merged:       {len(merged)}")

    mx.save_safetensors(str(out / "model.safetensors"),
                        merged, metadata={"format": "mlx"})
    print("Saved model.safetensors")

    build_config(dir_a, dir_b, out)
    copy_aux(dir_a, dir_b, out, template_from=args.template_from)

    print("\nDone. Test:")
    print(f"  python -m mlx_vlm.generate --model {out} "
          f"--prompt 'Hi' --max-tokens 30")
    print(f"  python -m mlx_vlm.generate --model {out} "
          f"--image test.jpg --prompt 'Describe this.' --max-tokens 100")


def build_config(dir_a: Path, dir_b: Path, out: Path):
    ca = json.loads((dir_a / "config.json").read_text())
    cb = json.loads((dir_b / "config.json").read_text())

    # A is authoritative for LM: it has the CORRECT quantization map
    # (A and B were quantized differently) and text_config (mtp=0).
    merged = dict(ca)

    # Graft vision-specific fields from B.
    for field in ("vision_config", "image_token_id", "video_token_id",
                  "vision_start_token_id", "vision_end_token_id"):
        if field in cb:
            merged[field] = cb[field]
            print(f"config: added '{field}' from B")

    (out / "config.json").write_text(json.dumps(merged, indent=2))
    print("config: wrote merged config.json "
          "(LM/quant from A, vision from B)")


def copy_aux(dir_a: Path, dir_b: Path, out: Path, template_from="a"):
    tmpl_dir = dir_b if template_from == "b" else dir_a
    print(f"aux: chat template / tokenizer_config from '{template_from}'")

    # Vocab files: identical between A and B; take from A.
    for n in ["tokenizer.json", "vocab.json", "merges.txt",
              "special_tokens_map.json", "added_tokens.json",
              "generation_config.json", "optiq_metadata.json"]:
        if (dir_a / n).exists():
            shutil.copy(dir_a / n, out / n)

    # Template-bearing files: from chosen source (default B, vision-aware).
    for n in ["tokenizer_config.json", "chat_template.jinja"]:
        if (tmpl_dir / n).exists():
            shutil.copy(tmpl_dir / n, out / n)

    # Vision preprocessing: always from B.
    for n in ["preprocessor_config.json", "processor_config.json",
              "image_processor_config.json", "video_processor_config.json"]:
        if (dir_b / n).exists():
            shutil.copy(dir_b / n, out / n)
            print(f"aux: copied {n} from B")


if __name__ == "__main__":
    main()