from pycsghub.upload_large_folder.main import upload_large_folder_internal
from pycsghub.upload_large_folder.main import create_repo
from pycsghub.csghub_api import CsgHubApi
import os
import argparse


BASE_URL = "http://10.20.100.5:4997"
TOKEN = "a2683e780fee4aa4b16cbedee0f75a45"

def upload(local_path, repo_id, repo_type="model", force=False):
    csg_api = CsgHubApi()
    use_full_repo_id = f"root/{repo_id}"

    repoExist, branchExist = csg_api.repo_branch_exists(
        repo_id=use_full_repo_id, repo_type=repo_type, revision="v1",
        endpoint=BASE_URL, token=TOKEN)
    
    if repoExist:
        if not force:
            print(f"{repo_id} has exists. Use --force to update it.")
            return
        else:
            print(f"{repo_id} has exists. Updating it because --force is set.")
    else:
        create_repo(
            api=csg_api,
            repo_id=use_full_repo_id,
            repo_type=repo_type,
            revision="v1",
            endpoint=BASE_URL,
            token=TOKEN
        )

    upload_large_folder_internal(
        repo_id=use_full_repo_id,
        local_path=local_path,
        repo_type=repo_type,
        revision="v1",
        endpoint=BASE_URL,
        token=TOKEN,
        allow_patterns=None,
        ignore_patterns=[".*"],  # 忽略所有隐藏目录和文件（以点开头的）
        num_workers=1,
        print_report=False,
        print_report_every=1,
    )


if __name__ == "__main__":
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description="自动上传模型到CSGHUB")
    
    # 添加本地模型路径参数
    parser.add_argument(
        "local_path", 
        type=str, 
        help="本地模型文件夹的路径"
    )
    
    # 添加强制更新参数
    parser.add_argument(
        "--force",
        action="store_true",
        help="强制更新远程仓库，即使仓库已存在"
    )
    
    # 解析命令行参数
    args = parser.parse_args()

    # 获取本地模型路径
    local_path = args.local_path
    
    # 判断本地模型路径是否存在
    if not os.path.exists(local_path):
        print(f"错误: 本地模型路径 '{local_path}' 不存在")
        exit(1)
    
    if not os.path.isdir(local_path):
        print(f"错误: 本地模型路径 '{local_path}' 不是一个文件夹")
        exit(1)
    
    # 从本地路径中提取仓库名称（使用文件夹名）
    repo_id = os.path.basename(local_path)
    
    print(f"开始上传本地模型: {local_path}")
    print(f"目标仓库: {repo_id}")
    
    # 执行上传
    upload(local_path, repo_id, force=args.force)