import os
import sys
from datasets import load_dataset, Dataset
from huggingface_hub import HfApi
import pandas as pd

TOKEN = os.environ.get("DEBUG")

api = HfApi(token=TOKEN)

out_dir = sys.argv[1]

#Uploading results
api.upload_folder(
    folder_path=out_dir,
    repo_id="EnergyStarAI/results_debug",
    repo_type="dataset",
)


#Updating requests
requests = load_dataset("EnergyStarAI/requests_debug", split="test", token=TOKEN)
requests_dset = requests.to_pandas()

models_ran=[]
for dir, path, files in os.walk(out_dir):
    for fid in files:
        model = os.path.join(dir, fid)
        models_ran.append(model)

print("Models ran are: " + models_ran)

requests_dset.loc[requests_dset["model"].isin(models_ran), ['status']] = "COMPLETED"
updated_dset =Dataset.from_pandas(requests_dset)
updated_dset.push_to_hub("EnergyStarAI/requests_debug", split="test", token=TOKEN)
print("Updated model status")