"""
Utility to safely inspect .safetensors checkpoints without opening binary files in text editors.
"""

import sys
import os
from safetensors import safe_open


def inspect_safetensors(filepath: str):
    if not os.path.exists(filepath):
        print(f"Error: File not found: {filepath}")
        return

    file_size_mb = os.path.getsize(filepath) / (1024 * 1024)
    print(f"==================================================")
    print(f" File: {os.path.basename(filepath)}")
    print(f" Path: {filepath}")
    print(f" Size: {file_size_mb:.2f} MB")
    print(f"==================================================")

    with safe_open(filepath, framework="pt") as f:
        meta = f.metadata() or {}
        print("\n--- Header Metadata ---")
        if meta:
            for k, v in meta.items():
                print(f"  {k}: {v}")
        else:
            print("  (No custom metadata embedded)")

        keys = list(f.keys())
        print(f"\n--- Tensors ({len(keys)} total) ---")
        for k in keys[:8]:
            t = f.get_tensor(k)
            print(f"  {k:45s} | shape: {str(list(t.shape)):20s} | dtype: {t.dtype}")

        if len(keys) > 8:
            print(f"  ... and {len(keys) - 8} more tensors.")


if __name__ == "__main__":
    target = sys.argv[1] if len(sys.argv) > 1 else "PRISM-1.safetensors"
    inspect_safetensors(target)