import os
import json

def collect_png_from_subfolders(base_path=".", output_file="../example/src/main/assets/pic/motion_data.json"):
    """
    遍历当前路径下所有的子文件夹（只遍历一层），收集PNG文件名并保存到JSON文件
    
    Args:
        base_path: 基础路径，默认为当前目录
        output_file: 输出的JSON文件名
    """
    # 检查基础路径是否存在
    if not os.path.exists(base_path):
        print(f"错误：路径 '{base_path}' 不存在")
        return {}
    
    # 获取所有子文件夹
    try:
        all_items = os.listdir(base_path)
    except PermissionError:
        print(f"错误：没有权限访问路径 '{base_path}'")
        return {}
    
    # 筛选出子文件夹（只遍历一层）
    subfolders = []
    for item in all_items:
        item_path = os.path.join(base_path, item)
        if os.path.isdir(item_path):
            subfolders.append(item)
    
    # 按文件夹名排序
    subfolders.sort()
    
    # 收集每个子文件夹中的PNG文件
    all_png_files = {}
    
    for folder in subfolders:
        folder_path = os.path.join(base_path, folder)
        
        try:
            folder_files = os.listdir(folder_path)
        except PermissionError:
            print(f"警告：没有权限访问文件夹 '{folder}'，跳过")
            continue
        
        # 筛选PNG文件
        png_filenames = []
        for file in folder_files:
            file_path = os.path.join(folder_path, file)
            if os.path.isfile(file_path) and file.lower().endswith('.png'):
                png_filenames.append(file)
        
        # 按文件名排序
        png_filenames.sort()
        
        # 如果有PNG文件，添加到结果中
        if png_filenames:
            all_png_files[folder] = png_filenames
    
    # 保存到JSON文件
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(all_png_files, f, ensure_ascii=False, indent=2)
        print(f"PNG文件信息已保存到 {output_file}")
    except Exception as e:
        print(f"保存文件时出错: {e}")
    
    # 显示统计信息
    total_folders = len(all_png_files)
    total_files = sum(len(files) for files in all_png_files.values())
    
    print(f"\n统计信息:")
    print(f"  扫描了 {len(subfolders)} 个子文件夹")
    print(f"  其中 {total_folders} 个文件夹包含PNG文件")
    print(f"  总共找到 {total_files} 个PNG文件")
    
    # 显示找到的文件
    if all_png_files:
        print("\n找到的PNG文件:")
        for folder, files in all_png_files.items():
            print(f"  {folder}/:")
            for filename in files:
                print(f"    - {filename}")
    
    return all_png_files

# 使用方法
if __name__ == "__main__":
    # 可以指定要遍历的基础路径，默认为当前目录
    base_path_to_scan = '../example/src/main/assets/pic'  # 当前目录
    
    collect_png_from_subfolders(base_path_to_scan)