from huggingface_hub import HfApi

from src.envs import LLM_CACHE_REPO, QUEUE_REPO, RESULTS_REPO, TOKEN


def check_and_create_dataset_repo(repo_id: str):
    api = HfApi(token=TOKEN)
    try:
        api.repo_info(repo_id=repo_id, repo_type="dataset")
        print(f"{repo_id} exists")
    except Exception:
        print(f"Creating {repo_id}")
        api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, private=True)


def check_and_create_repos():
    print("1. QUEUE Repository")
    check_and_create_dataset_repo(QUEUE_REPO)
    print("2. RESULTS Repository")
    check_and_create_dataset_repo(RESULTS_REPO)
    print("3. LLM Cache Repository")
    check_and_create_dataset_repo(LLM_CACHE_REPO)


if __name__ == "__main__":
    check_and_create_repos()