import os
import argparse
from safetensors import safe_open

def inspect_safetensors(file_path):
    if not os.path.exists(file_path):
        print(f"❌ 错误: 找不到文件 '{file_path}'")
        return

    print(f"🔍 正在读取文件: {file_path} ...")
    
    # 获取输出日志文件名
    base_name = os.path.splitext(os.path.basename(file_path))[0]
    output_txt = f"{base_name}_keys_list.txt"

    try:
        with safe_open(file_path, framework="pt", device="cpu") as f:
            keys = sorted(list(f.keys()))
            total_keys = len(keys)
            
            print(f"📊 文件读取成功！共有 {total_keys} 个权重键名 (Keys)。")
            print("-" * 50)
            print("💡 前 20 个键名预览:")
            
            # 打印前20个预览
            for i, key in enumerate(keys[:20]):
                tensor_slice = f.get_tensor(key)
                print(f"  [{i+1}] {key}  -->  shape: {list(tensor_slice.shape)} | dtype: {tensor_slice.dtype}")
            
            if total_keys > 20:
                print(f"  ... 以及其余 {total_keys - 20} 个键名")
            
            # 将所有键名导出到 txt 文件
            with open(output_txt, "w", encoding="utf-8") as txt_f:
                txt_f.write(f"模型文件: {file_path}\n")
                txt_f.write(f"总计键名数: {total_keys}\n")
                txt_f.write("=" * 60 + "\n")
                for key in keys:
                    tensor_slice = f.get_tensor(key)
                    txt_f.write(f"{key}\t{list(tensor_slice.shape)}\t{tensor_slice.dtype}\n")
            
            print("-" * 50)
            print(f"💾 完整键名列表已导出至当前目录: \n👉 {output_txt}")

    except Exception as e:
        print(f"❌ 读取过程中发生错误: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="快速检查 SafeTensors 模型内部的全部 Key 名")
    parser.add_argument("file_path", type=str, nargs="?", help="输入的 .safetensors 文件路径")
    args = parser.parse_args()

    # 如果运行脚本时没有传参，则自动在当前目录下找第一个 safetensors 文件
    if not args.file_path:
        current_dir_files = [f for f in os.listdir(".") if f.endswith(".safetensors")]
        if current_dir_files:
            print(f"提示: 未指定输入文件，自动选择当前目录下找到的: {current_dir_files[0]}")
            inspect_safetensors(current_dir_files[0])
        else:
            print("❌ 错误: 当前目录下没有找到任何 .safetensors 文件，请在后面加上文件路径参数。")
    else:
        inspect_safetensors(args.file_path)