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

templates_file = "GPTFuzzer.csv"
urllib.request.urlretrieve(
    "https://raw.githubusercontent.com/sherdencooper/GPTFuzz/refs/heads/master/datasets/prompts/GPTFuzzer.csv",
    templates_file)
out_file = "configuration/datasets/prompts/prompt_list.csv"

print(f"Reading prompt templates from GPTFuzzer: {templates_file}")
templates = []
out_data = []
for _, row in pd.read_csv(templates_file).iterrows():
    templates.append(row)

for id, entry in enumerate(templates):
    name = f"O{id}"
    text = entry["text"]

    out_entry = {}
    out_entry["id"] = id
    out_entry["template"] = name
    out_entry["text"] = text
    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(templates_file)
except OSError:
    pass
