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="AIEnergyScore/results_debug",
    repo_type="dataset",
)

# Updating requests
requests = load_dataset("AIEnergyScore/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("AIEnergyScore/requests_debug", split="test",
                         token=TOKEN)
print("Updated model status")