from huggingface_hub import HfApi

from src.envs import QUEUE_REPO, RESULTS_REPO, TOKEN


def check_and_create_repos():
    api = HfApi(token=TOKEN)

    # Check and create queue repo
    try:
        api.repo_info(repo_id=QUEUE_REPO, repo_type="dataset")
        print(f"Queue repository {QUEUE_REPO} exists")
    except Exception:
        print(f"Creating queue repository {QUEUE_REPO}")
        api.create_repo(repo_id=QUEUE_REPO, repo_type="dataset", exist_ok=True, private=False)

    # Check and create results repo
    try:
        api.repo_info(repo_id=RESULTS_REPO, repo_type="dataset")
        print(f"Results repository {RESULTS_REPO} exists")
    except Exception:
        print(f"Creating results repository {RESULTS_REPO}")
        api.create_repo(repo_id=RESULTS_REPO, repo_type="dataset", exist_ok=True, private=False)


if __name__ == "__main__":
    check_and_create_repos()