from huggingface_hub import HfApi
from huggingface_hub.utils import disable_progress_bars
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
import argparse
import os

disable_progress_bars()

parser = argparse.ArgumentParser()
parser.add_argument("-p", type=int, default=8, help="max parallel uploads")
args = parser.parse_args()

api = HfApi()
folder = Path(".")
repo_id = "NeutrinoLiu/rotbot_autoeval_ci"

def iter_files(base: Path, relative_base: Path):
    real_base = Path(os.path.realpath(base))
    for entry in real_base.iterdir():
        relative = relative_base / entry.name
        if entry.is_dir():
            yield from iter_files(entry, relative)
        elif entry.is_file():
            yield entry, relative

def upload_one(real: Path, relative: Path):
    api.upload_file(
        path_or_fileobj=str(real),
        path_in_repo=str(relative),
        repo_id=repo_id,
        repo_type="model",
    )
    return relative

files = list(iter_files(folder, Path(".")))

with ThreadPoolExecutor(max_workers=args.p) as pool:
    futures = {pool.submit(upload_one, real, rel): rel for real, rel in files}
    with tqdm(total=len(futures), desc="Uploading") as pbar:
        for future in as_completed(futures):
            rel = future.result()
            pbar.set_postfix_str(str(rel))
            pbar.update(1)