"""Inspect a LoRA .safetensors file: embedded training metadata + inferred
rank/target modules from tensor shapes. Read-only, no GPU/torch needed --
just parses the safetensors header, so it's fast and safe to run on any
LoRA file (including ones you didn't train yourself).

Usage:
    python_embeded\\python.exe dataset_builder\\inspect_lora.py path\\to\\lora.safetensors
"""
import argparse
import json
import re
from collections import Counter

from safetensors import safe_open


def parse_args():
    p = argparse.ArgumentParser(description="Inspect a LoRA safetensors file's metadata")
    p.add_argument("lora_path")
    return p.parse_args()


def infer_rank(keys, shapes):
    for k in keys:
        if k.endswith("lora_A.weight") or k.endswith("lora_down.weight"):
            return shapes[k][0]
    return None


def summarize_target_modules(keys):
    """Collapses per-block keys (blocks.0, blocks.1, ...) into module families
    so you can see what parts of the model were trained without a 800-line dump."""
    families = Counter()
    for k in keys:
        base = re.sub(r"\.(lora_A|lora_B|lora_down|lora_up|alpha)(\.weight)?$", "", k)
        base = re.sub(r"\.\d+\.", ".N.", base)
        families[base] += 1
    return families


def main():
    args = parse_args()
    with safe_open(args.lora_path, framework="pt") as f:
        meta = f.metadata() or {}
        keys = list(f.keys())
        shapes = {k: f.get_slice(k).get_shape() for k in keys}

    print(f"=== {args.lora_path} ===")
    print(f"Tensor count: {len(keys)}")
    print(f"Inferred rank (lora_A/lora_down first dim): {infer_rank(keys, shapes)}")

    print("\n--- Embedded metadata ---")
    if not meta:
        print("(none)")
    for k, v in meta.items():
        try:
            parsed = json.loads(v)
            print(f"{k}:")
            print(json.dumps(parsed, indent=2, ensure_ascii=False)[:3000])
        except (json.JSONDecodeError, TypeError):
            print(f"{k}: {v}")

    print("\n--- Target module families (per-block keys collapsed) ---")
    for base, count in summarize_target_modules(keys).most_common(30):
        print(f"  {base}  x{count}")


if __name__ == "__main__":
    main()