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 f in os.scandir(out_dir):
  if f.is_dir():
    for s in os.scandir(f):
      if s.is_dir() and s.name not in ['hooks','info','objects','refs','logs']:
          for m in os.scandir(s):
            models_ran.append(s.name+'/' + m.name)
              
print("Models ran are: " + str(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")