import os
import csv
from typing import Dict, List, Tuple, Any
from Tools.Calculate_Response_ASR import *


def read_prompts_csv(file_path: str) -> List[Tuple[str, int, int, str, str, str]]:
    """Read prompt datas from csv file, returns list of (dataset, category_id, task_id, category_name, question, instruction) tuples."""
    datas = []
    with open(file_path, mode="r", newline="", encoding="utf-8") as file:
      reader = csv.reader(file)
      datas = list(reader)[1:]
    return [[s[0], s[1], s[2], s[3], s[4].replace("\\n", "\n"), s[5].replace("\\n", "\n")] for s in datas]


def save_prompt_judgements_csv(judgements: List[Tuple[str, int, int, str, str, str, any]], output_file: str, judger_name: str = None, judgement_type: JudgementType = JudgementType.Origin, malicious_rate_str: str = None):
    """Save prompt jailbreak judgements to csv file."""
    os.makedirs(os.path.dirname(output_file), exist_ok=True)

    new_judgements = []
    for s in judgements:
        if judgement_type == JudgementType.Origin:
            j = s[6]
        elif judgement_type == JudgementType.Num_1:
            j = 1 if s[6].startswith("unsafe") else 0
        elif judgement_type == JudgementType.Num_10:
            j = 10 if s[6].startswith("unsafe") else 0
        elif judgement_type == JudgementType.Bool:
            j = True if s[6].startswith("unsafe") else False
        else:
            raise ValueError("Invalid judgement type!")

        new_judgements.append([s[0], s[1], s[2], s[3], s[4].replace("\n", "\\n"), s[5].replace("\n", "\\n"), j])


    # Construct header
    if malicious_rate_str is None:
      judge_title = f"Jailbreak Judgements - {judger_name}"
    else:
      judge_title = f"Jailbreak Judgements - {judger_name} ({malicious_rate_str})"
    new_judgements = [("dataset", "category_id", "task_id", "category_name", "question", "instruction", judge_title)] + new_judgements

    with open(output_file, mode="w", newline="", encoding="utf-8") as f:
      writer = csv.writer(f)
      writer.writerows(new_judgements)