#!/usr/bin/env python3
"""Write config.json + copy the auxiliary files for the re-packed checkpoint.

The weight re-packer only emits safetensors and an index. This adds everything else a
consumer needs, and -- importantly -- writes a quantization_config that DESCRIBES what we
actually did, because no released engine reads `deepseek_v41` yet and we are therefore
defining the convention rather than matching one.

Changes vs upstream config.json:
  expert_dtype        "fp4"  -> "nvfp4"     experts are now block-16 E2M1 + E4M3 scale
                                            + FP32 per-tensor global scale
  expert_block_size   (new)  16
  expert_scale_fmt    (new)  "e4m3"
  engram_dtype        (new)  "fp4"          Engram tables are now E2M1
  engram_block_size   (new)  32             keeping the source's own block/scale layout
  engram_scale_fmt    (new)  "ue8m0"
Everything else (dense fp8 block 32x32 / ue8m0) is untouched and left exactly as shipped.
"""
import glob, json, os, shutil, struct, sys

src, out = sys.argv[1], sys.argv[2]


def rebuild_index(out_dir):
    """Rebuild model.safetensors.index.json by reading every output shard's header.

    The re-packer is run in more than one pass (experts, then the two 94.5 GiB Engram
    shards), and each pass writes an index covering only its own shards -- so the last
    pass would otherwise clobber the rest. Reading the headers back is also a real check
    that every shard on disk is a valid safetensors file.
    """
    DT = {"I8": 1, "U8": 1, "F8_E4M3": 1, "F8_E8M0": 1, "BF16": 2, "F16": 2,
          "F32": 4, "F64": 8, "I32": 4, "I64": 8, "BOOL": 1}
    wmap, total = {}, 0
    shards = sorted(glob.glob(os.path.join(out_dir, "*.safetensors")))
    for path in shards:
        with open(path, "rb") as fh:
            n = struct.unpack("<Q", fh.read(8))[0]
            hdr = json.loads(fh.read(n))
        base = os.path.basename(path)
        for k, v in hdr.items():
            if k == "__metadata__":
                continue
            assert k not in wmap, f"tensor {k} appears in two shards"
            wmap[k] = base
            cnt = 1
            for d in v["shape"]:
                cnt *= d
            total += cnt * DT[v["dtype"]]
    return {"metadata": {"total_size": total}, "weight_map": wmap}, len(shards), total


idx, n_shards, total = rebuild_index(out)
json.dump(idx, open(os.path.join(out, "model.safetensors.index.json"), "w"), indent=1)
print(f"rebuilt index from {n_shards} shards: {len(idx['weight_map']):,} tensors, "
      f"{total/2**30:.1f} GiB")

cfg = json.load(open(os.path.join(src, "config.json")))
q = dict(cfg.get("quantization_config", {}))
assert q.get("expert_dtype") == "fp4", f"unexpected upstream quantization_config: {q}"
q.update({
    "expert_dtype": "nvfp4",
    "expert_block_size": 16,
    "expert_scale_fmt": "e4m3",
    "expert_global_scale": True,
    "engram_dtype": "fp4",
    "engram_block_size": 32,
    "engram_scale_fmt": "ue8m0",
    "repacked_by": "LibertAI/dsv41_fp4_stream.py",
})
cfg["quantization_config"] = q
json.dump(cfg, open(os.path.join(out, "config.json"), "w"), indent=2)
print("wrote config.json:", json.dumps(q))

FILES = ["tokenizer.json", "tokenizer_config.json", "LICENSE",
         "DeepSeek_V41_Tech_Report.pdf"]
DIRS = ["encoding", "inference"]
for fn in FILES:
    p = os.path.join(src, fn)
    if os.path.exists(p):
        shutil.copy2(p, os.path.join(out, fn)); print("  copied", fn)
for d in DIRS:
    p = os.path.join(src, d)
    if os.path.isdir(p):
        shutil.copytree(p, os.path.join(out, d), dirs_exist_ok=True); print("  copied", d + "/")

# sanity: index must cover every tensor the source had, minus the scale tensors we folded
si = json.load(open(os.path.join(src, "model.safetensors.index.json")))["weight_map"]
oi = json.load(open(os.path.join(out, "model.safetensors.index.json")))["weight_map"]
print(f"\nsource tensors {len(si):,} -> output tensors {len(oi):,}")
missing = [k for k in si if k not in oi and not k.endswith(".scale")]
extra = [k for k in oi if k not in si and not k.endswith((".weight_scale", ".weight_scale_2"))]
print(f"missing (non-scale): {len(missing)}  {missing[:3]}")
print(f"unexpected new:      {len(extra)}  {extra[:3]}")
assert not missing, "output is missing tensors the source had"