import random
from typing import Tuple

from turbofuzzllm.fuzzer.mutators import Mutator
from turbofuzzllm.fuzzer.template import Template
from turbofuzzllm.utils.logging import LOGGER
from turbofuzzllm.utils.prompt import QUESTION_PLACEHOLDER


class MutatePolicy:
    def __init__(self, mutators: "list[Mutator]", fuzzer: "TurboFuzzLLMFuzzer" = None):
        self.mutators = mutators
        self._fuzzer = fuzzer

    def select(self, template: "Template") -> Tuple["Mutator", str]:
        raise NotImplementedError("MutatePolicy must implement select method.")

    def update(self, mutant: "Template"):
        if QUESTION_PLACEHOLDER not in mutant.text:
            mutant.mutator.num_invalid_mutants += 1
        else:
            if mutant.num_jailbreaks > 0:
                mutant.mutator.num_jailbreaking_mutants += 1
            mutant.mutator.num_attacks += mutant.num_attacks
            mutant.mutator.num_rejects += mutant.num_rejects

    def mutate_single(self, template: "Template", mutator: "Mutator") -> "Template":
        mutated_text = mutator.mutate_template(template)
        mutator.num_mutants += 1
        name = f"G{Template.num_generated}"
        Template.num_generated += 1
        new_template = Template(
            text=mutated_text,
            name=name,
            parent=template,
            mutator=mutator,
            unexplored_child_mutations=self.mutators.copy(),
        )

        for m in self.mutators:
            if not m.can_mutate(new_template):
                new_template.unexplored_child_mutations.remove(m)

        return new_template

    @property
    def fuzzer(self):
        return self._fuzzer

    @fuzzer.setter
    def fuzzer(self, fuzzer):
        self._fuzzer = fuzzer
        for mutator in self.mutators:
            mutator.fuzzer = fuzzer


class MutateRandomSinglePolicy(MutatePolicy):
    def __init__(self, mutators: "list[Mutator]", fuzzer: "TurboFuzzLLMFuzzer" = None):
        super().__init__(mutators, fuzzer)

    def select(self, template: "Template") -> Tuple["Mutator", str]:
        counter = 0
        while True:
            mutator = random.choice(self.mutators)
            counter += 1
            if mutator.can_mutate(template) and mutator in template.unexplored_child_mutations:
                template.unexplored_child_mutations.remove(mutator)
                break
            if counter > 50:
                return mutator, "Failed"
        return mutator, "Succeeded"


class MutateRoundRobinSinglePolicy(MutatePolicy):
    def __init__(self, mutators: "list[Mutator]", fuzzer: "TurboFuzzLLMFuzzer" = None):
        super().__init__(mutators, fuzzer)
        self.index: int = 0

    def select(self, template: "Template") -> Tuple["Mutator", str]:
        while True:
            mutator = self.mutators[self.index]
            self.index = (self.index - 1 + len(self.mutators)) % len(self.mutators)
            if mutator.can_mutate(template):
                break
        return mutator, "Succeeded"


class MutateQLearningSinglePolicy(MutatePolicy):
    def __init__(self, mutators: "list[Mutator]", fuzzer: "TurboFuzzLLMFuzzer" = None):
        super().__init__(mutators, fuzzer)
        self.q_table: "dict[str, dict[Mutator, float]]" = dict()
        self.curr_action = None
        self.curr_options = None

        # Define hyper parameters
        self.learning_rate = 0.4
        self.discount_factor = 0.1
        self.epsilon = 1.0
        self.epsilon_decay_factor = 0.9
        self.epsilon_min = 0.2
        self.init_val = 0.0

    def log_q_table(self):
        LOGGER.debug(f"\t== MUTATION SELECTOR Q-LEARNING: Q Table ({len(self.q_table)} entries) ==")
        for state, action_values in self.q_table.items():
            LOGGER.debug(
                f"\t\t{state} -> {sorted(action_values.items(), key=lambda x: x[1], reverse=True)}"
            )
        LOGGER.debug(f"\t==\n")

    def calculate_reward(self, gen_template: "Template"):
        num_success = gen_template.num_jailbreaks
        num_failure = gen_template.num_rejects
        if QUESTION_PLACEHOLDER not in gen_template.text:
            num_success = 0
            num_failure = len(self.fuzzer.all_questions)
        return ((100.0 * num_success) / (num_success + num_failure)) - (
            (100.0 * num_failure) / (num_success + num_failure)
        )

    def select(self, template: "Template") -> Tuple["Mutator", str]:
        self.log_q_table()

        state = template.root
        if state not in self.q_table:
            self.q_table[state] = {m: self.init_val for m in self.mutators}
        state_table = self.q_table.get(state)

        self.curr_options = [
            m
            for m in self.mutators
            if m.can_mutate(template) and m in template.unexplored_child_mutations
        ]
        if len(self.curr_options) == 0:
            return self.mutators[0], "Failed"
        LOGGER.debug(f"\t== MUTATION SELECTOR Q-LEARNING: Valid mutations for {state} ==")
        for m in sorted(self.curr_options, key=lambda k: state_table.get(k), reverse=True):
            LOGGER.debug(f"\t\t{m} -> {state_table.get(m)}")
        LOGGER.debug(f"\t==\n")

        # Choose mutator with epsilon-greedy strategy
        if random.uniform(0, 1) < self.epsilon:
            # Explore
            self.curr_action = random.choice(self.curr_options)
            LOGGER.debug(f"\tMUTATION SELECTOR Q-LEARNING SELECT: Exploring {self.curr_action}")
        else:
            # Exploit
            offset = min(state_table.values())
            weights = [state_table.get(k) - offset + 1.0 for k in self.curr_options]
            self.curr_action = random.choices(self.curr_options, weights=weights)[0]
            LOGGER.debug(f"\tMUTATION SELECTOR Q-LEARNING SELECT: Exploiting {self.curr_action}")

        # decay epsilon
        self.epsilon = max(self.epsilon_min, self.epsilon_decay_factor * self.epsilon)
        template.unexplored_child_mutations.remove(self.curr_action)
        return self.curr_action, "Succeeded"

    def update(self, gen_template: "Template"):
        super().update(gen_template)
        state = gen_template.root
        reward = self.calculate_reward(gen_template)
        orig = self.q_table[state][self.curr_action]
        new = (1 - self.learning_rate) * orig + self.learning_rate * (
            reward
            + self.discount_factor * max([self.q_table[state].get(k) for k in self.curr_options])
        )
        LOGGER.debug(
            f"\tMUTATION SELECTOR Q-LEARNING UPDATE: Reward {reward} changed q value for ({state}, {self.curr_action}) from {orig} to {new}"
        )
        if (reward <= 0) and (new > orig):
            LOGGER.debug(f"WARN: negative reward increased q value")
            LOGGER.debug(f"ORIG: {orig}")
            LOGGER.debug(f"REWARD: {reward}")
            LOGGER.debug(f"MAX: {max(self.q_table[state].values())}")
            LOGGER.debug(f"PART1: {(1 - self.learning_rate) * orig}")
            LOGGER.debug(f"PART2: {self.learning_rate * reward}")
            LOGGER.debug(
                f"PART3: {self.learning_rate * self.discount_factor * max([self.q_table[state].get(k) for k in self.curr_options])}"
            )
            LOGGER.debug(f"NEW: {new}")
        self.q_table[state][self.curr_action] = new
