import os
import argparse
from huggingface_hub import HfApi

REPO_ID = "AIEnergyScore/benchmark-hugs-models"
TOKEN = os.environ.get("HF_TOKEN")
API = HfApi(token=TOKEN)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--hardware",
        default="a10g-large",
        type=str,
        required=False,
        help="hardware to use for benchmarking.",
    )
    args = parser.parse_args()
    return args
# Based on huggingface_hub _space_api.py
# CPU_BASIC = "cpu-basic"
#     CPU_UPGRADE = "cpu-upgrade"
#     T4_SMALL = "t4-small"
#     T4_MEDIUM = "t4-medium"
#     L4X1 = "l4x1"
#     L4X4 = "l4x4"
#     ZERO_A10G = "zero-a10g"
#     A10G_SMALL = "a10g-small"
#     A10G_LARGE = "a10g-large"
#     A10G_LARGEX2 = "a10g-largex2"
#     A10G_LARGEX4 = "a10g-largex4"
#     A100_LARGE = "a100-large"
#     V5E_1X1 = "v5e-1x1"
#     V5E_2X2 = "v5e-2x2"
#     V5E_2X4 = "v5e-2x4"
#curr_runtime = API.get_space_runtime(repo_id=REPO_ID)
#print(curr_runtime)
#requested_hardware = curr_runtime.requested_hardware
#print(requested_hardware)
#hardware_idx = hardware_options.index(requested_hardware)
#next_hardware = hardware_options[hardware_idx + 1]

if __name__ == '__main__':
    args = parse_args()
    curr_runtime = API.get_space_runtime(repo_id=REPO_ID)
    curr_hardware = curr_runtime.requested_hardware
    if curr_hardware != args.hardware:
        API.request_space_hardware(repo_id=REPO_ID, hardware=args.hardware)