import argparse
import os

from datasets import load_dataset, Dataset
from huggingface_hub import HfApi


TOKEN = os.environ.get("DEBUG")
api = HfApi(token=TOKEN)

parser = argparse.ArgumentParser()
parser.add_argument(
    "--run_dir",
    default=None,
    type=str,
    required=True,
    help="Path to the run directory.",
)
parser.add_argument(
    "--model_name",
    default=None,
    type=str,
    required=True,
    help="Model to benchmark.",
)
parser.add_argument(
    "--reason",
    default=None,
    type=str,
    required=False,
    help="Reason for failure -- to update in the requests file",
)

args = parser.parse_args()

# Updating request
dataset = load_dataset("AIEnergyScore/requests_debug", split="test", token=TOKEN).to_pandas()

## Set benchmark to failed

# If we have a custom reason for failure, add that instead of generic FAILED.
if args.reason:
    dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = args.reason
else:
    # TODO: This doesn't have to be try-except, we could actually check if the file is there...
    try:
        # Read error message
        with open(f"{args.run_dir}/error.log", 'r') as file:
            for f in file.readlines():
                if 'Traceback (most recent call last):' in f:
                    error_message = f
                    dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = "FAILED"
                    print("Status set to FAILED")
                else:
                    dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = "COMPLETED"
        # Add a new column for the error message if necessary
    except FileNotFoundError as e:
        print(f"Could not find {args.run_dir}/error.log")
    
updated_dataset = Dataset.from_pandas(dataset)
updated_dataset.push_to_hub("AIEnergyScore/requests_debug", split="test", token=TOKEN)