#!/usr/bin/env python
"""Merge LoRA into the QAT foundation, then quantize to compressed-tensors W4A16
(int4, group 32, symmetric) - Google's own recipe - for SGLang serving.

CRITICAL (handover 6.7): save_pretrained drops config keys the dataclass doesn't
declare. Backfill vision_config.model_patch_size / num_soft_tokens /
audio_config.output_proj_dims from the base config and ASSERT after every write.
"""
import os, sys, json, shutil, argparse, torch
os.environ.setdefault("HF_HOME","/workspace/.hf_home")

BASE_DEFAULT="/workspace/models/gemma4-qat-unquantized"
KEYS=(("vision_config","model_patch_size"),("vision_config","num_soft_tokens"),
      ("audio_config","output_proj_dims"))

def patch_cfg(out, base):
    """Backfill the three keys SGLang needs. Catches absent AND null."""
    cfg=os.path.join(out,"config.json"); b=json.load(open(os.path.join(base,"config.json")))
    c=json.load(open(cfg)); fixed=[]
    for sec,k in KEYS:
        c.setdefault(sec,{})
        if c[sec].get(k) is None:
            c[sec][k]=b[sec][k]; fixed.append(f"{sec}.{k}={b[sec][k]}")
    json.dump(c,open(cfg,"w"),indent=2)
    c=json.load(open(cfg))
    for sec,k in KEYS:
        assert c[sec].get(k) is not None, f"config lacks {sec}.{k} after patch"
    return fixed

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--adapter",required=True); ap.add_argument("--out",required=True)
    ap.add_argument("--base",default=BASE_DEFAULT)
    ap.add_argument("--merged-dir",default=None)
    ap.add_argument("--skip-quant",action="store_true")
    ap.add_argument("--quant-embeddings",action="store_true")
    A=ap.parse_args()
    merged=A.merged_dir or (A.out.rstrip("/")+"-bf16")
    from transformers import AutoTokenizer, AutoModelForCausalLM
    from peft import PeftModel
    print(f"[merge] base={A.base}\n[merge] adapter={A.adapter}",flush=True)
    tok=AutoTokenizer.from_pretrained(A.base)
    m=AutoModelForCausalLM.from_pretrained(A.base,dtype=torch.bfloat16,device_map="cpu")
    m=PeftModel.from_pretrained(m,A.adapter)
    m=m.merge_and_unload()                       # W' = W + BA, same shapes, zero added params
    os.makedirs(merged,exist_ok=True)
    m.save_pretrained(merged,safe_serialization=True)
    tok.save_pretrained(merged)
    for f in ("chat_template.jinja","processor_config.json","generation_config.json"):
        s=os.path.join(A.base,f)
        if os.path.exists(s): shutil.copy(s,os.path.join(merged,f))
    print("[merge] config backfill:",patch_cfg(merged,A.base),flush=True)
    n=sum(p.numel() for p in m.parameters())
    print(f"[merge] merged bf16 -> {merged}  params={n/1e9:.2f}B",flush=True)
    del m
    if A.skip_quant: return
    print("[quant] handing off to the isolated quant env (llm-compressor)",flush=True)

if __name__=="__main__": main()