import os
from huggingface_hub import HfApi

# ==========================================
# 配置区
# ==========================================
api = HfApi(token="***")

repo_id = "waveshaker/character_loras" # 你的仓库名
repo_type = "dataset" # 如果你的仓库是模型仓库，请改为 "model"
# ==========================================

def main():
    current_dir = os.getcwd()
    
    # 1. 交互输入地区前缀 (已按最新要求重排序)
    print("="*50)
    print("请选择这批数据上传到哪个地区前缀下：")
    print("1. CN | 2. KR | 3. JP | 4. Fictional | 5. WE | 6. 自定义")
    choice = input("\n👉 请输入数字或直接输入前缀 (默认回车为 KR): ").strip().upper()
    
    # 默认回车选 KR (2)
    if not choice: 
        choice = "2"
    
    # 映射字典同步更新
    mapping = {"1": "CN", "2": "KR", "3": "JP", "4": "Fictional", "5": "WE"}
    region_prefix = mapping.get(choice, choice).replace(" ", "_")
    
    # 2. 自动获取当前所在的母文件夹名字
    local_folder_name = os.path.basename(current_dir).replace(" ", "_")
    
    print(f"\n✅ 地区前缀: {region_prefix}")
    print(f"✅ 母文件夹名: {local_folder_name}")
    print("="*50 + "\n")

    # 3. 遍历并上传所有子文件夹
    for sub_folder_name in os.listdir(current_dir):
        sub_folder_path = os.path.join(current_dir, sub_folder_name)
        
        # 过滤掉非文件夹和隐藏文件
        if os.path.isdir(sub_folder_path) and not sub_folder_name.startswith('.'):
            
            # 动态构建云端路径: 例如 KR/YuJinNyang_KR2/1_clothing
            path_in_repo = f"{region_prefix}/{local_folder_name}/{sub_folder_name}"
            
            print(f"📦 开始上传: {sub_folder_name}")
            print(f"📂 目标路径: {path_in_repo}")
            
            try:
                api.upload_folder(
                    folder_path=sub_folder_path,
                    repo_id=repo_id,
                    repo_type=repo_type,
                    path_in_repo=path_in_repo,
                    # 动态 commit 信息
                    commit_message=f"Update files for {path_in_repo}",
                    # 👇 核心白名单：只上传这四种格式，完美过滤掉 .py 等垃圾文件
                    allow_patterns=["*.jpg", "*.jpeg", "*.txt", "*.png", "*.parquet"] 
                )
                print(f"🎉 {sub_folder_name} 上传完成！\n")
            except Exception as e:
                print(f"❌ {sub_folder_name} 上传失败: {e}\n")

    print("✅ 所有子文件夹处理完毕！")

if __name__ == "__main__":
    main()