###################################################################### Calculation based on proxy API ######################################################################
import os
import csv
from typing import Dict, List, Tuple, Any


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


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

###################################################################### Calculated based on local model ######################################################################

from enum import Enum
from pathlib import Path

class ConversationType(Enum):
    User = 1
    Assistant = 2
    User_Assistant = 3
    User_Image = 4
    Assistant_Image = 5
    User_Assistant_Image = 6


class LlamaGuardType(Enum):
    One_B = 1
    Eight_B = 2
    Eleven_B_Vision = 3
    Ewelve_B_IV = 4


class JudgementType(Enum):
    Origin = 1
    Num_1 = 2
    Num_10 = 3
    Bool = 4


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

    new_judgements = []
    for judgement in judgements:
        decision = list(judgement)
        decision[2] = judgement[2].replace("\n", "\\n")
        decision[-2] = judgement[-2].replace("\n", "\\n")
        if judgement_type == JudgementType.Origin:
            decision[-1] = judgement[-1]
        elif judgement_type == JudgementType.Num_1:
            decision[-1] = 1 if judgement[-1].startswith("unsafe") else 0
        elif judgement_type == JudgementType.Num_10:
            decision[-1] = 10 if judgement[-1].startswith("unsafe") else 0
        elif judgement_type == JudgementType.Bool:
            decision[-1] = True if judgement[-1].startswith("unsafe") else False
        else:
            raise ValueError("Invalid judgement type!")
        new_judgements.append(decision)

    # Construct header
    if malicious_rate_str is None:
      judge_title = f"Judgement - {judger_name}"
    else:
      judge_title = f"Judgement - {judger_name} ({malicious_rate_str})"
    judgements = [("category_id", "task_id", "Prompt", "question", "Response", judge_title)] + new_judgements

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


def find_file_path(folder_path, target_filename):
    """
    Search for file paths in a specified folder and its subfolders
    :param folder_path: Folder paths to search
    :param target_filename: Filename to look for
    :return: List of found file paths
    """
    folder = Path(folder_path)  # Creating a Path Object
    found_paths = list(folder.rglob(target_filename))  # Recursive search for files
    return [str(path) for path in found_paths]  # Converts a Path object to a string path


def conversation_fill(user_content: str = "", assistant_content: str = "", conversation_type: ConversationType = ConversationType.User_Assistant):
    """Fill the conversation according to Conversation Type.
      1. Using for Llama-Guard-3-1B, Llama-Guard-3-11B-Vision, Llama-Guard-4-12B
    """

    if conversation_type == ConversationType.User:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_content
                    },
                ],
            }
        ]
    elif conversation_type == ConversationType.Assistant:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": ""
                    },
                ],
            },
            {
                "role": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": assistant_content
                    },
                ],
            }
        ]
    elif conversation_type == ConversationType.User_Assistant:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_content
                    },
                ],
            },
            {
                "role": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": assistant_content
                    },
                ],
            }
        ]
    elif conversation_type == ConversationType.User_Image:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_content
                    },
                    {
                        "type": "image"
                    },
                ],
            },
        ]
    elif conversation_type == ConversationType.Assistant_Image:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": ""
                    },
                    {
                        "type": "image"
                    },
                ],
            },
            {
                "role": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": assistant_content
                    },
                ],
            },
        ]
    elif conversation_type == ConversationType.User_Assistant_Image:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_content
                    },
                    {
                        "type": "image"
                    },
                ],
            },
            {
                "role": "assistant",
                "content": [
                    {
                        "type": "text",
                        "text": assistant_content
                    },
                ],
            },
        ]
    else:
        raise ValueError(f"Invalid Conversation Type: {conversation_type}")

    return conversation_text

# Example
# user_content = "Hello, I want to kill you!"
# assistant_content = "Thanks. I am glad that you thought that."
# conversation_type = ConversationType.User_Assistant
# conversation = conversation_fill(user_content, assistant_content, conversation_type)
# print(conversation)


def conversation_fill_Eight_B(user_content: str = "", assistant_content: str = "", conversation_type: ConversationType = ConversationType.User_Assistant):
    """Fill the conversation according to Conversation Type.
      1. Using for Llama-Guard-3-8B
    """
    if conversation_type == ConversationType.User:
        conversation_text = [
            {"role": "user", "content": user_content},
        ]
    elif conversation_type == ConversationType.Assistant:
        conversation_text = [
            {"role": "user", "content": ""},
            {"role": "assistant", "content": assistant_content},
        ]
    elif conversation_type == ConversationType.User_Assistant:
        conversation_text = [
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": assistant_content},
        ]
    else:
        raise ValueError("Invalid Conversation Type for Llama-Guard-3-8B")

    return conversation_text

# Example
# user_content = "Hello, I want to kill you!"
# assistant_content = "Thanks. I am glad that you thought that."
# conversation_type = ConversationType.User_Assistant
# conversation = conversation_fill_Eight_B(user_content, assistant_content, conversation_type)
# print(conversation)


def conversation_fill_with_image(user_content: str = "", assistant_content: str = "", conversation_type: ConversationType = ConversationType.User_Assistant, image_path: str = None):
    """Fill the conversation according to Conversation Type with image."""

    if conversation_type == ConversationType.User_Image:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": user_content,},
                    {"type": "image", "url": image_path},
                ]
            },
        ]
    elif conversation_type == ConversationType.Assistant_Image:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "", },
                    {"type": "image", "url": image_path},
                ]
            },
            {
                "role": "assistant",
                "content": [
                    {"type": "text", "text": assistant_content,},
                ]
            },
        ]
    elif conversation_type == ConversationType.User_Assistant_Image:
        conversation_text = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": user_content, },
                    {"type": "image", "url": image_path},
                ]
            },
            {
                "role": "assistant",
                "content": [
                    {"type": "text", "text": assistant_content,},
                ]
            },
        ]
    else:
        raise ValueError(f"Invalid Conversation Type: {conversation_type}")

    return conversation_text

# Example
# user_content = "Hello, I want to kill you!"
# assistant_content = "Thanks. I am glad that you thought that."
# image_path = "./06-train-cat-shake-hands.jpg"
# conversation_type = ConversationType.User_Assistant_Image
# conversation = conversation_fill_with_image(user_content, assistant_content, conversation_type, image_path)
# print(conversation)