#!/usr/bin/env python3
"""Merge the ERP LoRA adapter into the bf16 base, producing servable weights.

WHY MERGE RATHER THAN HOT-SWAP. Serving NVFP4 base + LoRA at runtime was a
silent no-op on vLLM 0.24.0 (#47639, proven quant-agnostic). Merging first
sidesteps it entirely: the quantizer then sees ordinary bf16 weights and the
served artifact needs no adapter machinery at all.

⚠⚠ CHAT TEMPLATE. The trainee base ships a STALE 365-line chat_template.jinja;
upstream's is 390 lines. The harness deliberately trained through the UPSTREAM
template (config key `chat_template_path`), so the merged model MUST ship that
same upstream template. Shipping the base's own template here would be
train/serve skew with no error — it presents as a tuning failure.

⚠ CPU merge. device_map=None keeps the 48 GiB on host RAM (566 GB total here)
so this can run while GPU0 is training. Do not use device_map="auto".

⚠ Loader class. This checkpoint is Gemma4ForConditionalGeneration (vision +
audio towers present). Loading it as a plain CausalLM is playbook §3.2 — a
silent weight-load failure.
"""
import argparse
import json
import shutil
import sys
from pathlib import Path

UPSTREAM_TEMPLATE = "/tank/aimodels/gemma4-26b-a4b-it-bf16/chat_template.jinja"


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True)
    ap.add_argument("--adapter", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--chat-template", default=UPSTREAM_TEMPLATE)
    a = ap.parse_args()

    out = Path(a.out)
    if out.exists() and any(out.iterdir()):
        print(f"REFUSING: {out} exists and is non-empty", file=sys.stderr)
        return 1

    import torch
    from transformers import AutoTokenizer, Gemma4ForConditionalGeneration
    from peft import PeftModel

    print(f"[merge] loading base on CPU: {a.base}", flush=True)
    model = Gemma4ForConditionalGeneration.from_pretrained(
        a.base, dtype=torch.bfloat16, device_map=None, trust_remote_code=True,
    )

    # Count LoRA-target params before/after as a merge-actually-happened check.
    print(f"[merge] applying adapter: {a.adapter}", flush=True)
    before = {n: p.detach().clone() for n, p in model.named_parameters()
              if n.endswith("self_attn.q_proj.weight")
              and ".language_model.layers.0." in n}

    model = PeftModel.from_pretrained(model, a.adapter, is_trainable=False)
    n_lora = sum(1 for n, _ in model.named_parameters() if "lora_" in n)
    print(f"[merge] adapter tensors seen: {n_lora}", flush=True)
    if n_lora == 0:
        print("REFUSING: adapter contributed 0 tensors", file=sys.stderr)
        return 2

    model = model.merge_and_unload()
    print("[merge] merged", flush=True)

    # ⚠ Prove the merge changed weights. A no-op merge is the failure mode that
    # ships a base model wearing the tune's name, and nothing else would catch it.
    changed = 0
    for n, p in model.named_parameters():
        if n in before:
            if not torch.equal(p.detach(), before[n]):
                changed += 1
    if changed == 0:
        print("REFUSING: merge produced BIT-IDENTICAL weights on sampled "
              "LoRA-target modules — the adapter was inert or did not apply",
              file=sys.stderr)
        return 3
    print(f"[merge] verified {changed}/{len(before)} sampled target(s) changed", flush=True)

    out.mkdir(parents=True, exist_ok=True)
    print(f"[merge] saving to {out}", flush=True)
    model.save_pretrained(out, safe_serialization=True)

    # Tokenizer straight from the base — never one that has been through
    # calibration (playbook §3.14).
    AutoTokenizer.from_pretrained(a.base, trust_remote_code=True).save_pretrained(out)

    # ⚠ Ship the UPSTREAM chat template, matching what training rendered.
    src = Path(a.chat_template)
    if not src.exists():
        print(f"REFUSING: chat template missing at {src}", file=sys.stderr)
        return 4
    shutil.copy2(src, out / "chat_template.jinja")
    n_lines = len(src.read_text().splitlines())
    print(f"[merge] chat_template.jinja <- {src} ({n_lines} lines)", flush=True)

    # ⚠ CARRY THE PROCESSOR FILES. This is a multimodal checkpoint, so vLLM
    # builds a feature extractor at startup and dies without them:
    #     OSError: Can't load feature extractor for '<model>'
    # `save_pretrained` on the merged model writes tokenizer files only, so
    # anything else the base ships as auxiliary config must be copied across.
    # Verified against the served nvfp4a16 artifact, which carries exactly this.
    for aux in ("processor_config.json", "preprocessor_config.json",
                "video_preprocessor_config.json", "special_tokens_map.json"):
        src_aux = Path(a.base) / aux
        if src_aux.exists() and not (out / aux).exists():
            shutil.copy2(src_aux, out / aux)
            print(f"[merge] carried {aux}", flush=True)

    # ⚠⚠ CONFIG SCHEMA DOWNGRADE. transformers 5.15 MIGRATES Gemma-4's
    # heterogeneous-attention config on save: it drops `global_head_dim` /
    # `num_global_key_value_heads` and writes a `per_layer_config` dict instead.
    # Older transformers (5.10, which is what the llmcompressor venv pins) does
    # not understand the new key and resolves `config.num_key_value_heads` to
    # None, dying with:
    #     TypeError: unsupported operand type(s) for //: 'int' and 'NoneType'
    # Every WORKING artifact on this box - the bf16 base, the served nvfp4 prod
    # seat, and the nvfp4a16 build - uses the OLD schema. Merging a LoRA changes
    # weights, not architecture, so the base's expression of the architecture is
    # the correct one to ship.
    cfg_path = out / "config.json"
    cfg = json.loads(cfg_path.read_text())
    base_cfg = json.loads((Path(a.base) / "config.json").read_text())
    ct, bt = cfg.get("text_config", cfg), base_cfg.get("text_config", base_cfg)
    if "per_layer_config" in ct:
        ct.pop("per_layer_config")
        for k in ("global_head_dim", "num_global_key_value_heads"):
            if k in bt:
                ct[k] = bt[k]
        cfg_path.write_text(json.dumps(cfg, indent=2) + "\n")
        print("[merge] config schema downgraded to match the base "
              "(per_layer_config -> global_head_dim/num_global_key_value_heads)",
              flush=True)
    ct2 = json.loads(cfg_path.read_text()).get("text_config", {})
    for k in ("global_head_dim", "num_global_key_value_heads"):
        if bt.get(k) is not None and ct2.get(k) != bt.get(k):
            print(f"REFUSING: {k} is {ct2.get(k)}, base says {bt.get(k)}",
                  file=sys.stderr)
            return 6

    tj = out / "tokenizer.json"
    if tj.exists() and json.loads(tj.read_text()).get("truncation"):
        print("REFUSING: shipped tokenizer carries a truncation cap", file=sys.stderr)
        return 5

    print(f"[merge] DONE -> {out}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())