"""Playbook §4.3 post-steps for a Gemma-4 NVFP4 output tree.

Steps 1 and 3 (MTP graft, `re:^mtp.*` re-injection) are N/A on Gemma-4 -- it ships
no MTP head at all, verified as 0 mtp tensors in both bf16 sources. That leaves:

  step 2  restore processor_config.json + preprocessor_config.json
  step 4  confirm the saved tokenizer.json has truncation: null

Step 4 is not a formality here. The A4B output was quantized WITH the calibration
dataset, and build_calib calls the fast tokenizer with truncation=True,
max_length=8192 -- which mutates the Rust backend in place, and save_pretrained
then bakes the cap into the shipped tokenizer.json. It is latent on the
transformers that wrote it and fatal on a newer one. The fix edits the one
`truncation` key rather than copying the source file wholesale, so nothing else in
a 32 MB tokenizer can quietly change underneath it.

Derivation of preprocessor_config.json is `processor_config.json["image_processor"]`
verbatim; that reproduces the 2026-08-21 known-good output byte for byte.

Idempotent, and reports per step whether it CHANGED or was already correct.
Run with --check to verify without writing.
"""
import argparse
import json
import shutil
from pathlib import Path

ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True, help="bf16 source tree")
ap.add_argument("--out", required=True, help="quantized output tree")
ap.add_argument("--check", action="store_true", help="report only, write nothing")
a = ap.parse_args()
src, out = Path(a.src), Path(a.out)
mode = "CHECK" if a.check else "APPLY"
print(f"[{mode}] src={src}\n[{mode}] out={out}\n")

rc = 0


def step(n, desc):
    print(f"-- step {n}: {desc}")


step(1, "MTP graft")
mtp = [k for k in json.loads((out / "config.json").read_text()).get(
    "quantization_config", {}).get("ignore", []) if "mtp" in k.lower()]
idx = out / "model.safetensors.index.json"
tensors = json.loads(idx.read_text())["weight_map"] if idx.exists() else {}
n_mtp = sum(1 for k in tensors if k.startswith("mtp"))
print(f"   N/A for Gemma-4 (no MTP head). mtp tensors in output index: {n_mtp}; "
      f"mtp entries in ignore list: {len(mtp)}")
if n_mtp:
    print("   *** unexpected mtp tensors -- step 3 would become live, investigate")
    rc = 1

step(2, "restore processor_config.json + preprocessor_config.json")
spc = src / "processor_config.json"
if not spc.exists():
    print(f"   *** source has no processor_config.json -- cannot restore")
    rc = 1
else:
    opc = out / "processor_config.json"
    if opc.exists() and opc.read_bytes() == spc.read_bytes():
        print("   processor_config.json  already present and identical to source")
    elif a.check:
        print(f"   processor_config.json  MISSING/differs -> would copy from source")
    else:
        shutil.copy2(spc, opc)
        print("   processor_config.json  CHANGED (copied from source)")

    want = json.dumps(dict(json.loads(spc.read_text())["image_processor"]), indent=1)
    opre = out / "preprocessor_config.json"
    if opre.exists() and opre.read_text() == want:
        print("   preprocessor_config.json already present and correct")
    elif a.check:
        print("   preprocessor_config.json MISSING/differs -> would derive from image_processor")
    else:
        opre.write_text(want)
        print("   preprocessor_config.json CHANGED (derived from processor_config"
              "['image_processor'])")

step(4, "confirm saved tokenizer.json has truncation: null")
tj = out / "tokenizer.json"
tok = json.loads(tj.read_text())
trunc = tok.get("truncation")
if trunc is None:
    print("   truncation is null -- clean")
else:
    print(f"   *** truncation BAKED IN: {trunc}")
    stok = json.loads((src / "tokenizer.json").read_text())
    others = [k for k in set(tok) | set(stok)
              if k != "truncation" and tok.get(k) != stok.get(k)]
    print(f"   other top-level keys differing from source: {others or 'none'}")
    if a.check:
        print("   would set truncation -> null")
        rc = 1
    else:
        bak = tj.with_name("tokenizer.json.bak-pre-truncfix")
        if not bak.exists():
            shutil.copy2(tj, bak)
            print(f"   backed up -> {bak.name}")
        tok["truncation"] = None
        tmp = tj.with_suffix(".json.tmp")
        tmp.write_text(json.dumps(tok, ensure_ascii=False, indent=2))
        tmp.replace(tj)
        print("   truncation CHANGED -> null")

print(f"\n[{mode}] done rc={rc}")
raise SystemExit(rc)