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]
all_attempts_read = open("attempts.txt", "r+").readlines()
failed_attempts_read = open("failed_attempts.txt", "r+").readlines()

# Uploading output to the results dataset.
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()

for line in all_attempts_read:
    experiment_name, model = line.strip().split(',')
    if line not in failed_attempts_read:
        requests_dset.loc[
            requests_dset["model"] == model, ['status']] = "COMPLETED"
    else:
        requests_dset.loc[
            requests_dset["model"] == model, ['status']] = "FAILED"

updated_dset = Dataset.from_pandas(requests_dset)
updated_dset.push_to_hub("AIEnergyScore/requests_debug", split="test",
                         token=TOKEN)
print("Updated model status")