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(
    "--logs_name",
    default=None,
    type=str,
    required=False,
    help="Location of space runtime error log -- note this is distinct from an optimum-benchmark log.",
)
args = parser.parse_args()

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

# Set benchmark to failed
dataset.loc[dataset["model"].isin(args.model_name), ['status']] = "FAILED"

try:
    # Read error message
    with open(f"{args.run_dir}/error.log", 'r') as file:
       error_message = file.read()
    # Add a new column for the error message if necessary
    if "error_message" not in dataset.columns:
        dataset["error_message"] = ""
    dataset.loc[dataset["model"].isin(args.model_name), ['error_message']] = error_message
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("EnergyStarAI/requests_debug", split="test", token=TOKEN)

print("Status set to FAILED")

if args.logs_name:
    print("Attempting to save space runtime error log at EnergyStarAI/error_logs")
    try:
        api.upload_file(
            path_or_fileobj=args.error_log,
            path_in_repo=args.error_log,
            repo_id="EnergyStarAI/error_logs",
            repo_type="dataset",
        )
    except Exception as e:
        print("That didn't work. Error:")
        print(e)