#!/usr/bin/env python3
"""Mandatory post-steps after quantizing Qwen3.8-27B via the wrapper class.

The wrapper-class save drops the MTP head and the vision preprocessor configs.
All three of these have bitten previous rounds:

 1. graft `model-mtp.safetensors` verbatim from the bf16 source and register its
    tensors in the output index (else no speculative decoding at all);
 2. restore preprocessor_config.json / processor_config.json /
    video_preprocessor_config.json (else the vision tower can't preprocess);
 3. VERIFY `re:^mtp.*` is in quantization_config.ignore -- if it is missing,
    vLLM loads the grafted bf16 MTP head as though it were quantized and it
    comes up uninitialised, giving 0% acceptance. This is THE bug that cost two
    prior rounds; it is verified here rather than assumed.
"""
import json, os, shutil, sys

def main():
    src, out = sys.argv[1], sys.argv[2]
    fail = []

    # --- 1. MTP graft ---------------------------------------------------------
    # Two source layouts exist in the wild and both must work:
    #   (a) a standalone `model-mtp.safetensors` -- how JonathanColetti ships its
    #       grafted head, so a plain file copy suffices;
    #   (b) mtp.* living inside a NUMBERED shard -- how MuXodious/absolute-heresy
    #       ships (model-00012-of-00012.safetensors), because it is an unmodified
    #       full checkpoint rather than a graft.
    # Handling only (a) leaves the output index pointing at a `model-mtp.safetensors`
    # that was never created: the checkpoint looks fine to a tensor count but every
    # mtp tensor is unresolvable at load. Extract instead of copy for (b).
    mtp_src = os.path.join(src, "model-mtp.safetensors")
    mtp_dst = os.path.join(out, "model-mtp.safetensors")
    if os.path.exists(mtp_dst):
        print("MTP shard already present in output")
    elif os.path.exists(mtp_src):
        print(f"copying MTP shard ({os.path.getsize(mtp_src)/1e9:.2f} GB) ...", flush=True)
        shutil.copy2(mtp_src, mtp_dst)
    else:
        # layout (b): materialise the standalone shard the index will reference
        idx_path = os.path.join(src, "model.safetensors.index.json")
        wm = json.load(open(idx_path))["weight_map"]
        shards = sorted({wm[k] for k in wm if k.startswith("mtp")})
        if not shards:
            fail.append(f"no mtp.* in {src} (neither model-mtp.safetensors nor any shard)")
        else:
            from safetensors import safe_open
            from safetensors.torch import save_file
            print(f"extracting mtp.* from {shards} -> model-mtp.safetensors ...", flush=True)
            tensors = {}
            for shard in shards:
                with safe_open(os.path.join(src, shard), framework="pt") as f:
                    for k in f.keys():
                        if k.startswith("mtp"):
                            tensors[k] = f.get_tensor(k)
            save_file(tensors, mtp_dst, metadata={"format": "pt"})
            print(f"  wrote {len(tensors)} tensors, "
                  f"{os.path.getsize(mtp_dst)/1e6:.1f} MB")

    src_idx = json.load(open(os.path.join(src, "model.safetensors.index.json")))
    mtp_keys = [k for k in src_idx["weight_map"] if k.startswith("mtp")]
    out_idx_p = os.path.join(out, "model.safetensors.index.json")

    # A quant that lands under ~23 GB fits in ONE shard, and llm-compressor then
    # writes a bare `model.safetensors` with NO index at all. Every step below
    # needs one, so build it here rather than failing.
    #
    # Read the safetensors HEADER directly -- the first 8 bytes are a
    # little-endian u64 header length, followed by that many bytes of JSON
    # keyed by tensor name. Do NOT use safe_open() for this: it mmaps the whole
    # shard and ENOMEMs on ZFS against a 22 GB file.
    #
    # This has now bitten THREE separate rounds (2026-08-15, -08-20, -08-21),
    # each time fixed by hand and never in the script. Fixed in the script.
    if not os.path.exists(out_idx_p):
        import struct
        weight_map, total = {}, 0
        for fn in sorted(f for f in os.listdir(out) if f.endswith(".safetensors")):
            path = os.path.join(out, fn)
            total += os.path.getsize(path)
            with open(path, "rb") as fh:
                n = struct.unpack("<Q", fh.read(8))[0]
                header = json.loads(fh.read(n))
            for key in header:
                if key != "__metadata__":
                    weight_map[key] = fn
        json.dump({"metadata": {"total_size": total}, "weight_map": weight_map},
                  open(out_idx_p, "w"), indent=2)
        print(f"BUILT missing output index from safetensors headers: "
              f"{len(weight_map)} tensors across "
              f"{len(set(weight_map.values()))} shard(s), {total/1e9:.1f} GB")

    out_idx = json.load(open(out_idx_p))
    added = 0
    for k in mtp_keys:
        if k not in out_idx["weight_map"]:
            out_idx["weight_map"][k] = "model-mtp.safetensors"
            added += 1
    if added:
        json.dump(out_idx, open(out_idx_p, "w"), indent=2)
    print(f"MTP tensors in source: {len(mtp_keys)}; added to output index: {added}; "
          f"now present: {sum(1 for k in out_idx['weight_map'] if k.startswith('mtp'))}")
    if len(mtp_keys) == 0:
        fail.append("source index had NO mtp tensors")

    # --- 2. preprocessor / processor configs ---------------------------------
    for fn in ("preprocessor_config.json", "processor_config.json",
               "video_preprocessor_config.json", "chat_template.jinja",
               "generation_config.json"):
        s = os.path.join(src, fn)
        d = os.path.join(out, fn)
        if os.path.exists(s) and not os.path.exists(d):
            shutil.copy2(s, d)
            print(f"restored {fn}")
        elif os.path.exists(d):
            print(f"{fn} already present")
        else:
            print(f"NOTE: {fn} absent in source, skipped")
    if not os.path.exists(os.path.join(out, "preprocessor_config.json")):
        fail.append("preprocessor_config.json missing from output (vision will break)")

    # --- 3. verify the mtp ignore --------------------------------------------
    cfg_p = os.path.join(out, "config.json")
    cfg = json.load(open(cfg_p))
    ig = cfg.get("quantization_config", {}).get("ignore", [])
    has = any("mtp" in x for x in ig)
    if not has:
        # llm-compressor PRUNES ignore entries that matched no module at quant
        # time. The wrapper class never loads the MTP head, so `re:^mtp.*`
        # matches nothing and silently vanishes from the saved config -- and
        # then vLLM treats the freshly grafted bf16 MTP head as quantized and
        # brings it up uninitialised (0% acceptance). Re-inject it here, AFTER
        # the graft. This is the two-rounds-lost bug; repair, then re-verify.
        ig.append("re:^mtp.*")
        cfg["quantization_config"]["ignore"] = ig
        json.dump(cfg, open(cfg_p, "w"), indent=2)
        print("REPAIRED: re-injected 're:^mtp.*' into quantization_config.ignore "
              "(llm-compressor pruned it -- it matched no module at quant time)")
        cfg = json.load(open(cfg_p))
        ig = cfg["quantization_config"]["ignore"]
        has = any("mtp" in x for x in ig)
    print(f"quantization_config.ignore has an mtp entry: {has}  "
          f"({[x for x in ig if 'mtp' in x]})")
    if not has:
        fail.append("re:^mtp.* NOT in ignore -- MTP would load uninitialised (0% acceptance)")

    # --- report ---------------------------------------------------------------
    print("\nformat:", cfg.get("quantization_config", {}).get("format"))
    print("config_groups:", list(cfg.get("quantization_config", {}).get("config_groups", {})))
    if fail:
        print("\nFAILED CHECKS:")
        for f in fail:
            print("  -", f)
        return 1
    print("\nall post-steps OK")
    return 0

if __name__ == "__main__":
    sys.exit(main())