"""Fetch unchanged native carrier assets and omit replaced routed NVFP4 shards."""
import hashlib
import json
import re
import struct
from pathlib import Path

from huggingface_hub import HfApi, snapshot_download

repo = "local-inference-lab/GLM-5.3-Flash-NVFP4"
revision = "520de24eabf507659eaef7c70f14fd584527facc"
out = Path("/root/cv2data/out")
dest = out / "native-carrier"
names = [
    *(f"model-hf-nonexpert-{i:05d}-of-00004.safetensors" for i in range(1, 5)),
    *(f"model-mtp-{i:05d}-of-00004.safetensors" for i in range(1, 5)),
    "model-inputscales.safetensors", "model.safetensors.index.json",
    "config.json", "generation_config.json", "hf_quant_config.json",
    "processor_config.json", "tokenizer.json", "tokenizer_config.json",
    "chat_template.jinja", "LICENSE", "README.md",
    "amax.safetensors", "amax_checkpoint.safetensors", "amax_checkpoint.json",
]
info = HfApi().model_info(repo, revision=revision, files_metadata=True)
assert info.sha == revision
remote = {f.rfilename: f for f in info.siblings}
assert set(names) <= set(remote)
snapshot_download(repo, revision=revision, local_dir=dest, allow_patterns=names, max_workers=8)

records = []
for name in names:
    f = dest / name
    entry = remote[name]
    assert f.stat().st_size == entry.size, name
    with f.open("rb") as stream:
        digest = hashlib.file_digest(stream, "sha256").hexdigest()
    if entry.lfs:
        assert digest == entry.lfs.sha256, name
    else:
        data = f.read_bytes()
        assert hashlib.sha1(b"blob " + str(len(data)).encode() + b"\0" + data).hexdigest() == entry.blob_id, name
    records.append({"path": name, "bytes": entry.size, "source_sha256": digest})

index = json.loads((dest / "model.safetensors.index.json").read_text())
original_map = index["weight_map"]
selected = {name for name in names if name.startswith("model-") and name.endswith(".safetensors")}
omitted = {name: file for name, file in original_map.items() if file not in selected}
routed = re.compile(r"model\.language_model\.layers\.(?:[3-9]|[1-3][0-9]|4[0-4])\.mlp\.experts\.")
assert omitted and all(routed.match(name) for name in omitted), "Non-routed weights would be lost"
kept = {name: file for name, file in original_map.items() if file in selected}
seen = {}
tensor_bytes = 0
for name in sorted(selected):
    with (dest / name).open("rb") as stream:
        header_size = struct.unpack("<Q", stream.read(8))[0]
        header = json.loads(stream.read(header_size))
    for key, tensor in header.items():
        if key == "__metadata__":
            continue
        assert key not in seen
        seen[key] = name
        tensor_bytes += tensor["data_offsets"][1] - tensor["data_offsets"][0]
assert seen == kept, "Selected index and carrier tensor headers differ"
# Keep the carrier config intact: the runtime selects TrellisMX through the
# overlay environment variable and replaces the routed NVFP4 allocations.
index["weight_map"] = kept
index.setdefault("metadata", {})["total_size"] = tensor_bytes
(dest / "model.safetensors.index.json").write_text(json.dumps(index, indent=2) + "\n")
for record in records:
    f = dest / record["path"]
    record["bytes"] = f.stat().st_size
    record["sha256"] = record["source_sha256"]
    if record["path"] == "model.safetensors.index.json":
        record["sha256"] = hashlib.sha256(f.read_bytes()).hexdigest()
receipt = {
    "repo": repo, "revision": revision, "source_verification": "PASS",
    "index_header_coverage": "PASS", "indexed_tensors": len(kept),
    "omitted_replaced_routed_tensors": len(omitted),
    "bytes": sum(record["bytes"] for record in records), "files": records,
    "layout": "carrier/ under checkpoint root; MODEL_ROOT points to carrier/",
    "native_execution_tested": False,
}
(out / "native-carrier-receipt.json").write_text(json.dumps(receipt, indent=2) + "\n")
(out / "status/native-carrier.ready").write_text(revision + "\n")
print(json.dumps({k: v for k, v in receipt.items() if k != "files"}), flush=True)