import pandas as pd
import csv
import urllib.request
import os

questions_file = "harmful-behaviors.csv"
urllib.request.urlretrieve(
    "https://huggingface.co/datasets/JailbreakBench/JBB-Behaviors/resolve/main/data/harmful-behaviors.csv",
    questions_file)
out_file = "configuration/datasets/questions/jailbreakbench/harmful-behaviors.csv"

print(f"Reading harmful questions from JailBreakBench: {questions_file}")
questions = []
out_data = []
for _, row in pd.read_csv(questions_file).iterrows():
    questions.append(row)

for _, entry in enumerate(questions):
    index = entry["Index"]
    text = entry["Goal"]

    out_entry = {}
    out_entry["index"] = index
    out_entry["text"] = text
    for k, v in entry.items():
        if k != "Index" and k != "Goal":
            out_entry[k] = v
    out_data.append(out_entry)

def write_list_of_dicts_to_csv(data, filename):
    """
    Writes a list of dictionaries to a CSV file.

    Args:
        data: A list of dictionaries.
        filename: The name of the CSV file to write to.
    """
    if not data:
        return  # Nothing to write

    fieldnames = data[0].keys()

    with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(data)

write_list_of_dicts_to_csv(out_data, out_file)

try:
    os.remove(questions_file)
except OSError:
    pass
