from turbofuzzllm.fuzzer.question import Question
from turbofuzzllm.fuzzer.template import Template
from turbofuzzllm.utils.prompt import QUESTION_PLACEHOLDER

mutation_colors = {
    "IDENTITY": "darkgreen",
    "INJECT-PREFIX": "orchid",
    "REFUSAL-SUPPRESSION": "teal",
    "REPHRASE": "violet",
    "GENERATE-SIMILAR": "violetred",
    "CROSSOVER": "gold",
    "EXPAND-BEFORE": "olive",
    "EXPAND-AFTER": "limegreen",
    "SHORTEN": "tan",
    "FEW-SHOTS": "sienna",
    "TRANSFER-MUTATION": "slategray",
}


def draw_template_tree(
    templates: "list[Template]", questions: "dict[int, Question]", file_name: str, print_qid: bool
):
    n = 20
    top_n = [t for t in templates if t.num_attacks > 0]
    top_n = sorted(top_n, key=lambda x: (x.num_jailbreaks, x.asr), reverse=True)
    if n > 0:
        top_n = top_n[:n]
    best_templates = dict()
    for idx, t in enumerate(top_n):
        if t.num_jailbreaks > 0:
            best_templates[t] = idx + 1

    fp = open(file_name, "w", buffering=1)

    fp.write("digraph templates {\n")
    fp.write("  node[fontsize=18];\n")
    fp.write('  edge[lblstyle = "above, sloped" fontsize=14];\n')
    fp.write("\n")

    # print nodes
    for template in templates:
        label = f"{template}\\n{template.num_jailbreaks}/{template.asr:.0f}%"
        if print_qid and len(template.jailbroken_questions) != 0:
            label += "\n"
            for idx, qid in enumerate(template.jailbroken_questions):
                label += f"{questions[qid]}"
                if idx != len(template.jailbroken_questions) - 1:
                    label += ", "
        shape = "oval" if template.root_parent == template else "box"
        color = f"0.000 {template.num_jailbreaks / 2.0 / len(questions):.3f} 1.000"
        if template in best_templates:
            label += f"\\n#{best_templates.get(template)}"
            if best_templates.get(template) <= 5:
                shape = "doubleoctagon"
            elif best_templates.get(template) <= 10:
                shape = "Msquare"
        fp.write(
            f'  {template} [label = "{label}" shape = "{shape}" style = "filled" fillcolor = "{color}"];\n'
        )
    fp.write("\n")

    # print edges
    for template in templates:
        for child in template.children:
            rhs = f"{child}"
            if QUESTION_PLACEHOLDER not in child.text:
                rhs = f'"{child}\\nX"'
            label = f"{str(child.mutator)[0] + str(child.mutator)[-1]}"
            color = mutation_colors[f"{child.mutator}"]
            fp.write(
                f'  {template} -> {rhs} [label = "{label}" fontcolor = "{color}" color = "{color}"];\n'
            )
    fp.write("\n")

    fp.write("}\n")
    fp.close()
