import warnings

from turbofuzzllm.llm.bedrock_judge import BedrockJudge
from turbofuzzllm.llm.gpt_judge import GPTJudge
from turbofuzzllm.llm.harmbench_judge import HarmBenchJudge
from turbofuzzllm.llm.model_utils import (
    is_bedrock_model,
    is_local_harmbench_model,
    is_openai_model,
    is_sagemaker_harmbench_model,
    is_sagemaker_jumpstart_llama_guard_model,
)
from turbofuzzllm.llm.roberta_judge import RoBERTaJudge
from turbofuzzllm.llm.sm_harmbench_judge import SagemakerHarmBenchJudge
from turbofuzzllm.llm.sm_llama_guard_judge import SagemakerLlamaGuardJudge
from turbofuzzllm.utils.cli_utils import setup_cli_args

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)

import argparse
import os
import random
import shutil
import sys
import time

import pandas as pd

current = os.path.dirname(os.path.realpath(__file__))
parent = os.path.dirname(current)
sys.path.append(parent)

from turbofuzzllm.fuzzer import (
    EXP3SelectPolicy,
    GPTMutatorCrossOver,
    GPTMutatorExpandAfter,
    GPTMutatorExpandBefore,
    GPTMutatorFewShots,
    GPTMutatorGenerateSimilar,
    GPTMutatorRephrase,
    GPTMutatorShorten,
    GPTMutatorTransferMutation,
    MCTSExploreSelectPolicy,
    MutateQLearningSinglePolicy,
    MutateRandomSinglePolicy,
    MutateRoundRobinSinglePolicy,
    QLearningSelectPolicy,
    RandomSelectPolicy,
    RoundRobinSelectPolicy,
    SimpleMutatorIdentity,
    SimpleMutatorInjectPrefix,
    SimpleMutatorRefusalSuppression,
    TurboFuzzLLMFuzzer,
    UCBSelectPolicy,
)
from turbofuzzllm.fuzzer.template import find_initial_asr, has_positive_initial_asr
from turbofuzzllm.llm import MutatorModel, TargetModel
from turbofuzzllm.utils import LOGGER, setup_logging, setup_stats, stats

__version__ = "1.5.0"
header = """
-----------------------------------------------------------------------------------------------------------------
TurboFuzzLLM: Turbocharging Mutation-based Fuzzing for Effectively Jailbreaking Large Language Models in Practice
-----------------------------------------------------------------------------------------------------------------
"""


def add_common_args(p):
    model_group = p.add_argument_group("target/attacker/judge models")
    model_group.add_argument(
        "--target-model-id",
        dest="target_model_id",
        type=str,
        default="us.anthropic.claude-3-5-haiku-20241022-v1:0",
        # default="us.anthropic.claude-3-5-sonnet-20241022-v2:0",
        # default="gpt-4o",
        # default="gpt-4o-mini",
        # default="gpt-4-turbo",
        # default="gpt-3.5-turbo",
        # default="gpt-3.5-turbo-1106",
        # default="gpt-4-0613",
        # default="gpt-4-1106-preview",
        # default="google/gemma-7b",
        # default="google/gemma-2b",
        # default="HuggingFaceH4/zephyr-7b-beta",
        # default ="cais/zephyr_7b_r2d2",
        # default="anthropic.claude-3-sonnet-20240229-v1:0",
        # default="anthropic.claude-3-haiku-20240307-v1:0",
        help="Target model id (default: %(default)s)",
    )
    model_group.add_argument(
        "--mutator-model-id",
        dest="mutator_model_id",
        type=str,
        default="mistral.mistral-large-2407-v1:0",
        help="Mutator model id (default: %(default)s)",
    )
    model_group.add_argument(
        "--judge-model-id",
        dest="judge_model_id",
        type=str,
        default="anthropic.claude-3-sonnet-20240229-v1:0",
        # default="cais/HarmBench-Llama-2-13b-cls",
        # default="hubert233/GPTFuzz",
        help="Judge model id (default: %(default)s)",
    )
    model_group.add_argument(
        "--judge-tokenizer",
        dest="judge_tokenizer",
        type=str,
        default="-",
        # default="cais/HarmBench-Llama-2-13b-cls",
        # default="hubert233/GPTFuzz",
        help="Judge tokenizer path (default: %(default)s)",
    )
    model_group.add_argument(
        "--guardrail-id",
        dest="guardrail_id",
        type=str,
        default="-",
        help="Target model guardrail id (default: %(default)s)",
    )
    model_group.add_argument(
        "--guardrail-version",
        dest="guardrail_version",
        type=str,
        default="-",
        help="Target model guardrail version (default: %(default)s)",
    )

    access_group = p.add_argument_group("aws access arguments")
    access_group.add_argument(
        "--region",
        dest="region",
        type=str,
        default="us-west-2",
        # default="us-east-1",
        help="AWS region (default: %(default)s)",
    )
    access_group.add_argument(
        "--sagemaker-role",
        dest="sagemaker_role",
        type=str,
        default="-",
        help="Sagemaker execution role ARN (default: %(default)s)",
    )
    access_group.add_argument(
        "--sagemaker-mutator-instance-type",
        dest="sagemaker_mutator_instance_type",
        type=str,
        default="ml.g5.12xlarge",
        help="Sagemaker mutator instance type (default: %(default)s)",
    )
    access_group.add_argument(
        "--sagemaker-target-instance-type",
        dest="sagemaker_target_instance_type",
        type=str,
        default="ml.g5.12xlarge",
        help="Sagemaker target instance type (default: %(default)s)",
    )
    access_group.add_argument(
        "--sagemaker-judge-instance-type",
        dest="sagemaker_judge_instance_type",
        type=str,
        default="ml.g5.12xlarge",
        help="Sagemaker judge instance type (default: %(default)s)",
    )
    access_group.add_argument(
        "--api-key",
        dest="api_key",
        type=str,
        default="-",
        help="API key (default: %(default)s)",
    )


def add_answer_args(parser_ans):
    answer_group = parser_ans.add_argument_group("answer mode arguments")
    answer_group.add_argument(
        "--max-queries",
        dest="max_queries",
        type=int,
        default=100,
        help="The maximum number of queries in total (default: %(default)s)",
    )
    answer_group.add_argument(
        "--template-selector",
        dest="template_selector",
        type=str,
        default="ql",
        help="Template selector: rand (random), rr (round-robin), ql (q-learning), ucb, "
        "mcts, exp3 (default: %(default)s)",
    )
    answer_group.add_argument(
        "--mutation-selector",
        dest="mutation_selector",
        type=str,
        default="ql",
        help="Mutation selector: rand (random), rr (round-robin), ql (q-learning) (default: %(default)s)",
    )
    answer_group.add_argument(
        "--early-stopping-constant",
        dest="early_stopping_constant",
        type=int,
        default=4,
        help="The maximum number of queries in total (default: %(default)s)",
    )
    answer_group.add_argument(
        "--warmup-percent",
        dest="warmup_percent",
        type=int,
        default=0,
        help="The percentage of queries to use during warmup (default: %(default)s)",
    )
    answer_group.add_argument(
        "--max-jailbreaks-per-question",
        dest="max_jailbreaks_per_question",
        type=int,
        default=1,
        help=argparse.SUPPRESS,
    )

    input_group = parser_ans.add_argument_group("input file paths")
    input_group.add_argument(
        "--templates-path",
        dest="templates_path",
        type=str,
        default="auto",
        help="Path to initial templates CSV (default: %(default)s)",
    )

    config_group = parser_ans.add_argument_group("common config arguments")
    config_group.add_argument(
        "--num-threads", dest="num_threads", type=int, default=1, help=argparse.SUPPRESS
    )
    config_group.add_argument(
        "--seed", dest="seed", type=int, default=0, help="Random seed (default: %(default)s)"
    )
    config_group.add_argument(
        "--run-name",
        dest="run_name",
        type=str,
        default="-",
        help="Suffix to append as run name (default: %(default)s)",
    )
    config_group.add_argument(
        "--max-templates",
        dest="max_templates",
        type=int,
        default=20,
        help="The maximum number of initial templates (default: %(default)s)",
    )
    config_group.add_argument(
        "--template", dest="template", type=str, help="Select a single template"
    )

    add_common_args(parser_ans)


def add_attack_args(parser_a):
    attack_group = parser_a.add_argument_group("attack mode arguments")
    attack_group.add_argument(
        "--max-queries",
        dest="max_queries",
        type=int,
        default=4000,
        help="The maximum number of queries in total (default: %(default)s)",
    )
    attack_group.add_argument(
        "--template-selector",
        dest="template_selector",
        type=str,
        default="ql",
        help="Template selector: rand (random), rr (round-robin), ql (q-learning), ucb, "
        "mcts, exp3 (default: %(default)s)",
    )
    attack_group.add_argument(
        "--mutation-selector",
        dest="mutation_selector",
        type=str,
        default="ql",
        help="Mutation selector: rand (random), rr (round-robin), ql (q-learning) (default: %(default)s)",
    )
    attack_group.add_argument(
        "--early-stopping-constant",
        dest="early_stopping_constant",
        type=int,
        default=4,
        help="The maximum number of queries in total (default: %(default)s)",
    )
    attack_group.add_argument(
        "--warmup-percent",
        dest="warmup_percent",
        type=int,
        default=100,
        help="The percentage of queries to use during warmup (default: %(default)s)",
    )
    attack_group.add_argument(
        "--max-jailbreaks-per-question",
        dest="max_jailbreaks_per_question",
        type=int,
        default=1,
        help=argparse.SUPPRESS,
    )

    input_group = parser_a.add_argument_group("input file paths")
    input_group.add_argument(
        "--templates-path",
        dest="templates_path",
        type=str,
        default="configuration/datasets/prompts/prompt_list.csv",
        help="Path to initial templates CSV (default: %(default)s)",
    )
    input_group.add_argument(
        "--questions-path",
        dest="questions_path",
        type=str,
        default="configuration/datasets/questions/harmbench/harmbench_behaviors_text_standard_all.csv",
        help="Path to questions CSV (default: %(default)s)",
    )

    config_group = parser_a.add_argument_group("common config arguments")
    config_group.add_argument(
        "--num-threads",
        dest="num_threads",
        type=int,
        default=1,
        help="Number of threads to use (default: %(default)s)",
    )
    config_group.add_argument(
        "--seed", dest="seed", type=int, default=0, help="Random seed (default: %(default)s)"
    )
    config_group.add_argument(
        "--run-name",
        dest="run_name",
        type=str,
        default="-",
        help="Suffix to append as run name (default: %(default)s)",
    )
    config_group.add_argument(
        "--max-templates",
        dest="max_templates",
        type=int,
        default=-1,
        help="The maximum number of initial templates (default: %(default)s)",
    )
    config_group.add_argument(
        "--max-questions",
        dest="max_questions",
        type=int,
        default=-1,
        help="The maximum number of questions (default: %(default)s)",
    )
    config_group.add_argument(
        "--question-id", dest="question_id", type=int, help="Select a single question"
    )
    config_group.add_argument(
        "--template", dest="template", type=str, help="Select a single template"
    )

    add_common_args(parser_a)


def add_legacy_args(parser_l):
    legacy_group = parser_l.add_argument_group("legacy mode arguments")
    legacy_group.add_argument(
        "--max-queries",
        dest="max_queries",
        type=int,
        default=4000,
        help="The maximum number of queries in total (default: %(default)s)",
    )
    legacy_group.add_argument(
        "--template-selector",
        dest="template_selector",
        type=str,
        default="mcts",
        help="Template selector: rand (random), rr (round-robin), ql (q-learning), ucb, "
        "mcts, exp3 (default: %(default)s)",
    )
    legacy_group.add_argument(
        "--mutation-selector",
        dest="mutation_selector",
        type=str,
        default="rand",
        help="Mutation selector: rand (random), rr (round-robin), ql (q-learning) (default: %(default)s)",
    )
    legacy_group.add_argument(
        "--early-stopping-constant",
        dest="early_stopping_constant",
        type=int,
        default=-1,
        help=argparse.SUPPRESS,
    )
    legacy_group.add_argument(
        "--warmup-percent",
        dest="warmup_percent",
        type=int,
        default=0,
        help=argparse.SUPPRESS,
    )
    legacy_group.add_argument(
        "--max-jailbreaks-per-question",
        dest="max_jailbreaks_per_question",
        type=int,
        default=1,
        help="The maximum jailbreaks per question (default: %(default)s)",
    )

    input_group = parser_l.add_argument_group("input file paths")
    input_group.add_argument(
        "--templates-path",
        dest="templates_path",
        type=str,
        default="configuration/datasets/prompts/prompt_list.csv",
        help="Path to initial templates CSV (default: %(default)s)",
    )
    input_group.add_argument(
        "--questions-path",
        dest="questions_path",
        type=str,
        default="configuration/datasets/questions/harmbench/harmbench_behaviors_text_standard_all.csv",
        help="Path to questions CSV (default: %(default)s)",
    )

    config_group = parser_l.add_argument_group("common config arguments")
    config_group.add_argument(
        "--num-threads",
        dest="num_threads",
        type=int,
        default=1,
        help="Number of threads to use (default: %(default)s)",
    )
    config_group.add_argument(
        "--seed", dest="seed", type=int, default=0, help="Random seed (default: %(default)s)"
    )
    config_group.add_argument(
        "--run-name",
        dest="run_name",
        type=str,
        default="-",
        help="Suffix to append as run name (default: %(default)s)",
    )
    config_group.add_argument(
        "--max-templates",
        dest="max_templates",
        type=int,
        default=-1,
        help="The maximum number of initial templates (default: %(default)s)",
    )
    config_group.add_argument(
        "--max-questions",
        dest="max_questions",
        type=int,
        default=-1,
        help="The maximum number of questions (default: %(default)s)",
    )
    config_group.add_argument(
        "--question-id", dest="question_id", type=int, help="Select a single question"
    )
    config_group.add_argument(
        "--template", dest="template", type=str, help="Select a single template"
    )

    add_common_args(parser_l)


def add_evaluate_args(parser_e):
    evaluate_group = parser_e.add_argument_group("evaluate mode arguments")
    evaluate_group.add_argument(
        "--max-queries",
        dest="max_queries",
        type=int,
        default=-1,
        help=argparse.SUPPRESS,
    )
    evaluate_group.add_argument(
        "--early-stopping-constant",
        dest="early_stopping_constant",
        type=int,
        default=-1,
        help=argparse.SUPPRESS,
    )
    evaluate_group.add_argument(
        "--warmup-percent",
        dest="warmup_percent",
        type=int,
        default=0,
        help=argparse.SUPPRESS,
    )
    evaluate_group.add_argument(
        "--max-jailbreaks-per-question",
        dest="max_jailbreaks_per_question",
        type=int,
        default=-1,
        help="The maximum jailbreaks per question (default: %(default)s)",
    )

    input_group = parser_e.add_argument_group("input file paths")
    input_group.add_argument(
        "--templates-path",
        dest="templates_path",
        type=str,
        default="configuration/datasets/prompts/prompt_list.csv",
        help="Path to initial templates CSV (default: %(default)s)",
    )
    input_group.add_argument(
        "--questions-path",
        dest="questions_path",
        type=str,
        default="configuration/datasets/questions/jailbreakbench/harmful-behaviors.csv",
        help="Path to questions CSV (default: %(default)s)",
    )

    config_group = parser_e.add_argument_group("common config arguments")
    config_group.add_argument(
        "--num-threads",
        dest="num_threads",
        type=int,
        default=1,
        help="Number of threads to use (default: %(default)s)",
    )
    config_group.add_argument(
        "--seed", dest="seed", type=int, default=0, help="Random seed (default: %(default)s)"
    )
    config_group.add_argument(
        "--run-name",
        dest="run_name",
        type=str,
        default="-",
        help="Suffix to append as run name (default: %(default)s)",
    )
    config_group.add_argument(
        "--max-templates",
        dest="max_templates",
        type=int,
        default=-1,
        help="The maximum number of initial templates (default: %(default)s)",
    )
    config_group.add_argument(
        "--max-questions",
        dest="max_questions",
        type=int,
        default=-1,
        help="The maximum number of questions (default: %(default)s)",
    )
    config_group.add_argument(
        "--question-id", dest="question_id", type=int, help="Select a single question"
    )
    config_group.add_argument(
        "--template", dest="template", type=str, help="Select a single template"
    )

    add_common_args(parser_e)


def get_parser():
    # create the top-level parser
    parser = argparse.ArgumentParser(prog="turbofuzzllm")
    parser.add_argument("-v", "--version", action="version", version="%(prog)s " + __version__)
    subparsers = parser.add_subparsers(
        dest="mode", help="Mode to run: answer, attack, evaluate, or legacy"
    )

    # create the parser for the "answer" command
    parser_ans = subparsers.add_parser(
        "answer", help="Jailbreak a single question provided interactively by user"
    )
    add_answer_args(parser_ans)

    # create the parser for the "attack" command
    parser_a = subparsers.add_parser(
        "attack", help="Jailbreak each question once as fast as possible"
    )
    add_attack_args(parser_a)

    # create the parser for the "evaluate" command
    parser_e = subparsers.add_parser(
        "evaluate", help="Evaluate input templates on each question to calculate ASR"
    )
    add_evaluate_args(parser_e)

    # create the parser for the "legacy" command
    parser_g = subparsers.add_parser(
        "legacy",
        help="Run legacy/gptfuzzer style search that evaluates a template with each question",
    )
    add_legacy_args(parser_g)

    return parser


def execute(cli_args):
    # setup output path
    date_key = time.strftime("%Y-%m-%d", time.localtime())
    time_key = time.strftime("%H-%M-%S", time.localtime())
    run_name = f"{cli_args.mode}_{cli_args.target_model_id.replace('/', '-')}"
    if cli_args.run_name != "-":
        run_name += f"_{cli_args.run_name}"
    out_path = f"output/{date_key}/{run_name}_{time_key}"
    if os.path.exists(out_path) and os.path.isdir(out_path):
        shutil.rmtree(out_path)
    os.makedirs(out_path)

    # setup logger and stats writer
    setup_logging(f"{out_path}/details.log")
    setup_stats(f"{out_path}/stats.txt")
    LOGGER.info(header)
    LOGGER.info(f"Output folder: {out_path}")
    LOGGER.info(f"Logging in file details.log")
    LOGGER.info(f"Recording statistics in file stats.txt")

    stats().write("mode", cli_args.mode)

    stats().write("random-seed", cli_args.seed)
    random.seed(cli_args.seed)

    stats().write("num-threads", cli_args.num_threads)

    if cli_args.templates_path == "auto":
        cli_args.templates_path = (
            f"configuration/datasets/prompts/learnt/templates-{cli_args.target_model_id}.csv"
        )
        if not os.path.isfile(cli_args.templates_path):
            cli_args.templates_path = "configuration/datasets/prompts/prompt_list.csv"
    if not os.path.exists(cli_args.templates_path):
        LOGGER.info(f"Unable to find original templates from file {cli_args.templates_path}")
        raise FileNotFoundError(
            f"Unable to find original templates from file {cli_args.templates_path}. Use scripts/get_templates_gptfuzzer.py to get original templates."
        )
    LOGGER.info(f"Reading input templates from file {cli_args.templates_path}")
    original_templates = dict()
    for _, row in pd.read_csv(cli_args.templates_path).iterrows():
        original_templates[row["template"]] = row
    if cli_args.template is not None:
        if cli_args.template not in original_templates:
            raise Exception(
                f"Template with name {cli_args.template} not found in {cli_args.templates_path}"
            )
        original_templates = {cli_args.template: original_templates[cli_args.template]}
    else:
        if cli_args.mode == "evaluate" and cli_args.max_templates == -1:
            count_positive_asr = 0
            for idx, data in original_templates.items():
                if has_positive_initial_asr(data):
                    count_positive_asr += 1
            cli_args.max_templates = count_positive_asr
        if (cli_args.max_templates > 0) and (len(original_templates) > cli_args.max_templates):
            sorted_indices = sorted(original_templates.keys(), key=lambda x: random.random())
            sorted_indices = sorted(
                sorted_indices, key=lambda x: find_initial_asr(original_templates[x]), reverse=True
            )
            sorted_indices = sorted_indices[: cli_args.max_templates]
            original_templates = {idx: original_templates[idx] for idx in sorted_indices}

    if cli_args.mode == "answer":
        questions = dict()
        question_text = input("Enter a harmful question to jailbreak and press ENTER: ")
        questions[0] = pd.Series({"index": 0, "text": question_text})
    else:
        LOGGER.info(f"Reading input questions from file {cli_args.questions_path}")
        questions = dict()
        for _, row in pd.read_csv(cli_args.questions_path).iterrows():
            questions[int(row["index"])] = row
        if cli_args.question_id is not None:
            if cli_args.question_id not in questions:
                raise Exception(
                    f"Question with id {cli_args.question_id} not found in {cli_args.questions_path}"
                )
            questions = {cli_args.question_id: questions[cli_args.question_id]}
        else:
            if (cli_args.max_questions > 0) and (len(questions) > cli_args.max_questions):
                sorted_indices = sorted(questions.keys(), key=lambda x: random.random())
                sorted_indices = sorted_indices[: cli_args.max_questions]
                questions = {idx: questions[idx] for idx in sorted_indices}

    select_policy = RoundRobinSelectPolicy()
    if cli_args.mode != "evaluate":
        if cli_args.template_selector == "ql":
            select_policy = QLearningSelectPolicy()
        elif cli_args.template_selector == "mcts":
            select_policy = MCTSExploreSelectPolicy()
        else:
            if cli_args.template_selector == "rand":
                select_policy = RandomSelectPolicy()
            elif cli_args.template_selector == "rr":
                select_policy = RoundRobinSelectPolicy()
            elif cli_args.template_selector == "ucb":
                select_policy = UCBSelectPolicy()
            elif cli_args.template_selector == "exp3":
                select_policy = EXP3SelectPolicy()
            else:
                raise Exception(f"Unexpected template selector mode: {cli_args.template_selector}")
    stats().write("template-select-policy", type(select_policy).__name__)

    stats().write("max-queries", cli_args.max_queries)

    max_warmup_queries = int((cli_args.max_queries * cli_args.warmup_percent) / 100)
    stats().write("warmup-percent", cli_args.warmup_percent)
    stats().write("max-warmup-queries", max_warmup_queries)

    stats().write("max-jailbreaks-per-question", cli_args.max_jailbreaks_per_question)

    mutator_model = None
    target_model = None
    judge_model = None
    try:
        mutator_model = MutatorModel(cli_args.mutator_model_id, cli_args.api_key)
        stats().write("mutator-model", cli_args.mutator_model_id)

        guardrails_enabled = (cli_args.guardrail_id != "-") and (cli_args.guardrail_version != "-")
        target_model = TargetModel(
            cli_args.target_model_id,
            cli_args.api_key,
            cli_args.guardrail_id,
            cli_args.guardrail_version,
        )
        stats().write("target-model", cli_args.target_model_id)
        stats().write("target-guardrail-id", cli_args.guardrail_id)
        stats().write("target-guardrail-version", cli_args.guardrail_version)

        if is_local_harmbench_model(cli_args.judge_model_id):
            judge_model = HarmBenchJudge(
                cli_args.judge_model_id, cli_args.api_key, cli_args.judge_tokenizer
            )
        elif is_sagemaker_harmbench_model(cli_args.judge_model_id):
            judge_model = SagemakerHarmBenchJudge(cli_args.judge_model_id, cli_args.api_key)
        elif is_openai_model(cli_args.judge_model_id):
            judge_model = GPTJudge(cli_args.judge_model_id, cli_args.api_key)
        elif is_sagemaker_jumpstart_llama_guard_model(cli_args.judge_model_id):
            assert cli_args.sagemaker_role != "-"
            judge_model = SagemakerLlamaGuardJudge(cli_args.judge_model_id, cli_args.api_key)
        elif "/" in cli_args.judge_model_id:
            judge_model = RoBERTaJudge(
                cli_args.judge_model_id, cli_args.api_key, cli_args.judge_tokenizer
            )
        else:
            assert is_bedrock_model(cli_args.judge_model_id)
            judge_model = BedrockJudge(cli_args.judge_model_id, cli_args.api_key)
        stats().write("judge-model", cli_args.judge_model_id)

        mutators = []

        if cli_args.mode == "attack" or cli_args.mode == "answer":
            mutators.append(SimpleMutatorIdentity())
            mutators.append(SimpleMutatorInjectPrefix())
            mutators.append(SimpleMutatorRefusalSuppression())
            mutators.append(GPTMutatorExpandAfter(mutator_model))
            mutators.append(GPTMutatorFewShots(mutator_model))
            mutators.append(GPTMutatorTransferMutation(mutator_model))

        mutators.append(GPTMutatorCrossOver(mutator_model))
        mutators.append(GPTMutatorRephrase(mutator_model))
        mutators.append(GPTMutatorGenerateSimilar(mutator_model))
        mutators.append(GPTMutatorShorten(mutator_model))
        mutators.append(GPTMutatorExpandBefore(mutator_model))

        mutate_policy = MutateRoundRobinSinglePolicy(mutators)
        if cli_args.mode != "evaluate":
            if cli_args.mutation_selector == "ql":
                mutate_policy = MutateQLearningSinglePolicy(mutators)
            elif cli_args.mutation_selector == "rand":
                mutate_policy = MutateRandomSinglePolicy(mutators)
            elif cli_args.mutation_selector == "rr":
                mutate_policy = MutateRoundRobinSinglePolicy(mutators)
            else:
                raise Exception(f"Unexpected mutation selector mode: {cli_args.mutation_selector}")
        stats().write("mutation-select-policy", type(mutate_policy).__name__)

        setup_cli_args(cli_args)

        fuzzer = TurboFuzzLLMFuzzer(
            questions=questions,
            attacker=mutator_model,
            target=target_model,
            judge=judge_model,
            initial_templates=original_templates,
            mutate_policy=mutate_policy,
            select_policy=select_policy,
            max_jailbreaks_per_question=cli_args.max_jailbreaks_per_question,
            max_queries=cli_args.max_queries,
            max_warmup_queries=max_warmup_queries,
            out_path=out_path,
            mode=cli_args.mode,
            num_threads=cli_args.num_threads,
            guardrails_enabled=guardrails_enabled,
            early_stopping_constant=cli_args.early_stopping_constant,
        )

        fuzzer.run()
    except Exception as e:
        if mutator_model is not None:
            mutator_model.cleanup()
        if target_model is not None:
            target_model.cleanup()
        if judge_model is not None:
            judge_model.cleanup()
        raise e


def turbofuzzllm():
    cli_args = get_parser().parse_args()
    setup_cli_args(cli_args)
    execute(cli_args)
