#!/usr/bin/env python3
"""
Verify safetensors file integrity: header matches actual data offsets.
Used for diagnosing 'Error while deserializing header: incomplete metadata'
even when file size looks correct.

Usage from Proxmox host (remote ComfyUI in LXC):
    ssh root@10.0.20.91 'bash -s' < scripts/check_safetensors.py
"""
import struct, json, os, sys

def check(path):
    if not os.path.exists(path):
        print(f"MISSING: {path}")
        return False
    try:
        with open(path, "rb") as f:
            hlen = struct.unpack("<Q", f.read(8))[0]
            header = json.loads(f.read(hlen))
            total = 8 + hlen
            for k, v in header.items():
                if isinstance(v, dict) and "data_offsets" in v:
                    total += v["data_offsets"][1] - v["data_offsets"][0]
            actual = os.path.getsize(path)
            ok = total == actual
            print(f"{os.path.basename(path)}: exp={total:,} act={actual:,} ok={ok}")
            return ok
    except Exception as e:
        print(f"{os.path.basename(path)}: ERROR: {type(e).__name__}: {e}")
        return False

if __name__ == "__main__":
    if len(sys.argv) < 2:
        # Default: check common ComfyUI model locations
        paths = [
            "/opt/ComfyUI/models/diffusion_models/ideogram4_fp8_scaled.safetensors",
            "/opt/ComfyUI/models/diffusion_models/ideogram4_unconditional_fp8_scaled.safetensors",
            "/opt/ComfyUI/models/text_encoders/qwen3vl_8b_fp8_scaled.safetensors",
            "/opt/ComfyUI/models/vae/flux2-vae.safetensors",
        ]
    else:
        paths = sys.argv[1:]
    
    results = [check(p) for p in paths]
    if not all(results):
        sys.exit(1)
    print("All OK")