"""Verify every downloaded weight: exact byte size vs ModelScope, and a safetensors header parse."""
import os, sys, json, struct, urllib.request

BASE = os.environ.get("DIFFSYNTH_MODEL_BASE_PATH", "/data/models/minimax-h3/models")
WANT = {
    "DiffSynth-Studio/MiniMax-H3-NF4": ["minimax-h3-fl2va-nf4.safetensors",
                                        "minimax-h3-ref2va-nf4.safetensors",
                                        "minimax-h3-text-encoder-nf4.safetensors",
                                        "video_vae_nf4.safetensors", "audio_vae_nf4.safetensors"],
    "lightx2v/Minimax-h3-Turbo":       ["minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors",
                                        "minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors"],
}
PROCESSORS = ["MiniMax/MiniMax-H3/FL2VA/processor", "MiniMax/MiniMax-H3/Ref2VA/processor"]

def remote_sizes(mid):
    u = f"https://www.modelscope.cn/api/v1/models/{mid}/repo/files?Recursive=true"
    d = json.load(urllib.request.urlopen(u, timeout=60))
    return {f["Path"]: f.get("Size", 0) for f in d.get("Data", {}).get("Files", []) if f.get("Type") == "blob"}

def header_ok(path):
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        if not (0 < n < 200_000_000): return False, f"bad header len {n}"
        hdr = json.loads(f.read(n))
        end = max(v["data_offsets"][1] for k, v in hdr.items() if k != "__metadata__")
        expect = 8 + n + end
        actual = os.path.getsize(path)
        return expect == actual, f"tensors={len(hdr)} expect={expect} actual={actual}"

bad = 0
for mid, files in WANT.items():
    rs = remote_sizes(mid)
    for fn in files:
        p = os.path.join(BASE, mid, fn)
        if not os.path.exists(p):
            print(f"MISSING  {mid}/{fn}"); bad += 1; continue
        got, want = os.path.getsize(p), rs.get(fn, -1)
        size_ok = (got == want)
        ok, detail = header_ok(p)
        flag = "OK  " if (size_ok and ok) else "BAD "
        if not (size_ok and ok): bad += 1
        print(f"{flag} {fn}\n       size {got} vs remote {want} {'==' if size_ok else '!= MISMATCH'} | header {'ok' if ok else 'BROKEN'} ({detail})")

need = ["tokenizer.json", "tokenizer_config.json", "preprocessor_config.json", "chat_template.json", "vocab.json", "merges.txt"]
for processor in PROCESSORS:
    pdir = os.path.join(BASE, processor)
    missing = [f for f in need if not os.path.exists(os.path.join(pdir, f))]
    print((f"OK   {processor}: all present" if not missing else f"BAD  {processor} missing: {missing}"))
    bad += len(missing)
print("\nRESULT:", "ALL GOOD" if bad == 0 else f"{bad} PROBLEM(S)")
sys.exit(1 if bad else 0)