import statistics
from dataclasses import dataclass
from typing import TYPE_CHECKING

from pandas import Series

if TYPE_CHECKING:
    from turbofuzzllm.fuzzer.mutators import Mutator


def find_initial_asr(data: Series) -> float:
    # try getting average asr for all asr columns
    all_asr = []
    for k, v in data.items():
        if str(k).startswith("asr"):
            all_asr.append(float(v))
    if len(all_asr) != 0:
        return statistics.fmean(all_asr)
    return 0.0


def has_positive_initial_asr(data: Series) -> float:
    init_asr = find_initial_asr(data)
    return init_asr > 0.0


@dataclass
class Response:
    text: str
    result: str


class Template:
    num_generated = 0

    def __init__(
        self,
        text: str,
        name: str,
        data: "Series" = None,
        parent: "Template" = None,
        mutator: "Mutator" = None,
        unexplored_child_mutations: "list[Mutator]" = None,
    ):

        self.text: str = text
        self.name: str = name
        self.data: "Series" = data

        self.children: "list[Template]" = []
        if parent is None:
            self.is_original: bool = True
            self.parent: "Template" = self
            self.root_parent: "Template" = self
            self.level: int = 0
            if not self.name.startswith("O"):
                self.name = "O" + self.name
        else:
            self.is_original = False
            self.parent: "Template" = parent
            self.parent.children.append(self)
            self.root_parent: "Template" = parent.root_parent
            self.level: int = parent.level + 1

        self.mutator: "Mutator" = mutator

        self.responses: "dict[int, Response]" = dict()
        self.jailbroken_questions: "set[int]" = set()
        self.visited_num = 0

        self.initial_asr = 0.0
        self.unexplored_child_mutations = unexplored_child_mutations

        if self.data is not None:
            self.initial_asr = find_initial_asr(data)
            pass
        else:
            assert self.root_parent is not None
            self.initial_asr = self.root_parent.initial_asr

    def __repr__(self):
        return self.name

    def __str__(self):
        return self.name

    @property
    def root(self):
        return f"{self.root_parent}"

    @property
    def question_indices(self):
        return self.responses.keys()

    @property
    def num_jailbreaks(self):
        return len(self.jailbroken_questions)

    @property
    def num_attacks(self):
        return len(self.responses)

    @property
    def num_rejects(self):
        return self.num_attacks - self.num_jailbreaks

    @property
    def asr(self):
        if len(self.responses) == 0:
            return 0.0
        return (100.0 * self.num_jailbreaks) / len(self.question_indices)
