import os
from fastapi import FastAPI
from huggingface_hub import HfApi
import time

TOKEN = os.environ.get("BULK_ENERGY_TOKEN")
API = HfApi(token=TOKEN)

REPO_ID = "AIEnergyScore/BulkCalcSpace"
RESULTS_DSET = "AIEnergyScore/BulkCalcResults"
app = FastAPI()

@app.get("/")
def start_train():
    model_file = open("models.txt", "r+").readlines()
    task_file = open("tasks.txt", "r+").readlines()
    hardware_file = open("hardware.txt", "r+").readlines()
    for hardware in hardware_file:
        hardware = hardware.strip()
        os.system(f"echo 'Requested hardware is {hardware}'")
        curr_runtime = API.get_space_runtime(repo_id=REPO_ID)
        os.system(f"echo 'Current hardware is {curr_runtime}'")
        if curr_runtime.hardware != hardware:
            os.system("echo 'Trying to switch.'")
            API.request_space_hardware(repo_id=REPO_ID, hardware=hardware)
        for model in model_file:
            model = model.strip()
            os.system(f"echo 'Attempting to benchmark model {model}.'")
            for task in task_file:
                task = task.strip()
                os.system(f"echo 'Attempting to benchmark model {model} on task {task}.'")
                # Create the name of the directory for output.
                now = time.time()
                run_dir = f"runs/{task}/{model}/{now}"
                os.system(f"./entrypoint.sh {REPO_ID} {model} {task} {hardware} {run_dir}")
                # Uploads all run output to the results dataset.
                os.system(f"echo 'Uploading {run_dir} to {RESULTS_DSET}'")
                try:
                    API.create_repo(repo_id=f"{RESULTS_DSET}", repo_type="dataset",)
                    os.system("echo 'Created results dataset repository.'")
                except:
                    os.system("echo 'Using pre-existing dataset respository.'")
                API.upload_folder(folder_path=run_dir, repo_id=f"{RESULTS_DSET}", repo_type="dataset",)
    print("Pausing space")
    API.pause_space(REPO_ID)
    #return {"Status": "Done"}