import math
import random
from typing import TYPE_CHECKING

import numpy as np

from turbofuzzllm.utils.logging import LOGGER
from turbofuzzllm.utils.priority_queue import PriorityQueue

if TYPE_CHECKING:
    from turbofuzzllm.fuzzer.core import TurboFuzzLLMFuzzer
    from turbofuzzllm.fuzzer.template import Template

from turbofuzzllm.utils.prompt import QUESTION_PLACEHOLDER


class SelectPolicy:
    def __init__(self, fuzzer: "TurboFuzzLLMFuzzer"):
        self.fuzzer = fuzzer
        self.template_policy_name = "SelectPolicy"

    def select(self) -> "Template":
        raise NotImplementedError("SelectPolicy must implement select method.")

    def update(self, template: "Template"):
        pass

    def initialize(self):
        pass


class RoundRobinSelectPolicy(SelectPolicy):
    def __init__(self, fuzzer: "TurboFuzzLLMFuzzer" = None):
        super().__init__(fuzzer)
        self.index: int = 0

    def select(self) -> "Template":
        seed = self.fuzzer.templates[self.index]
        seed.visited_num += 1
        return seed

    def update(self, template: "Template"):
        self.index = (self.index - 1 + len(self.fuzzer.templates)) % len(self.fuzzer.templates)


class RandomSelectPolicy(SelectPolicy):
    def __init__(self, fuzzer: "TurboFuzzLLMFuzzer" = None):
        super().__init__(fuzzer)

    def select(self) -> "Template":
        seed = random.choice(self.fuzzer.templates)
        seed.visited_num += 1
        return seed


class UCBSelectPolicy(SelectPolicy):
    def __init__(self, explore_coeff: float = 1.0, fuzzer: "TurboFuzzLLMFuzzer" = None):
        super().__init__(fuzzer)

        self.step = 0
        self.last_choice_index = None
        self.explore_coeff = explore_coeff
        self.rewards = [0 for _ in range(len(self.fuzzer.templates))]

    def select(self) -> "Template":
        if len(self.fuzzer.templates) > len(self.rewards):
            self.rewards.extend([0 for _ in range(len(self.fuzzer.templates) - len(self.rewards))])

        self.step += 1
        scores = np.zeros(len(self.fuzzer.templates))
        for i, template in enumerate(self.fuzzer.templates):
            smooth_visited_num = template.visited_num + 1
            scores[i] = self.rewards[i] / smooth_visited_num + self.explore_coeff * np.sqrt(
                2 * np.log(self.step) / smooth_visited_num
            )

        self.last_choice_index = np.argmax(scores)
        self.fuzzer.templates[self.last_choice_index].visited_num += 1
        return self.fuzzer.templates[self.last_choice_index]

    def update(self, template: "Template"):
        succ_num = template.num_jailbreaks
        self.rewards[self.last_choice_index] += succ_num / len(self.fuzzer.all_questions)


class MCTSExploreSelectPolicy(SelectPolicy):
    def __init__(self, fuzzer: "TurboFuzzLLMFuzzer" = None, ratio=0.5, alpha=0.1, beta=0.2):
        super().__init__(fuzzer)

        self.step = 0
        self.mctc_select_path: "list[Template]" = []
        self.last_choice = None
        self.rewards: "dict[Template, int]" = dict()
        self.ratio = ratio  # balance between exploration and exploitation
        self.alpha = alpha  # penalty for level
        self.beta = beta  # minimal reward after penalty

        self.template_policy_name = "mcts"

    def initialize(self):
        for pnode in self.fuzzer.initial_templates:
            if pnode not in self.rewards:
                self.rewards[pnode] = pnode.initial_asr

    def select(self) -> "Template":
        self.step += 1

        self.mctc_select_path.clear()
        cur = max(
            self.fuzzer.initial_templates,
            key=lambda pnode: self.rewards.get(pnode, 0) / (pnode.visited_num + 1)
            + self.ratio * np.sqrt(2 * np.log(self.step) / (pnode.visited_num + 0.01)),
        )
        self.mctc_select_path.append(cur)

        while len(cur.children) > 0:
            if np.random.rand() < self.alpha:
                break
            cur = max(
                cur.children,
                key=lambda pnode: self.rewards.get(pnode, 0) / (pnode.visited_num + 1)
                + self.ratio * np.sqrt(2 * np.log(self.step) / (pnode.visited_num + 0.01)),
            )
            self.mctc_select_path.append(cur)

        for pn in self.mctc_select_path:
            pn.visited_num += 1

        self.last_choice = cur
        return cur

    def update(self, template: "Template"):
        succ_num = template.num_jailbreaks

        for pn in reversed(self.mctc_select_path):
            reward = succ_num / len(self.fuzzer.all_questions)
            if pn not in self.rewards:
                self.rewards[pn] = 0
            self.rewards[pn] += reward * max(self.beta, (1 - 0.1 * self.last_choice.level))


class EXP3SelectPolicy(SelectPolicy):
    def __init__(self, gamma: float = 0.05, alpha: float = 25, fuzzer: "TurboFuzzLLMFuzzer" = None):
        super().__init__(fuzzer)

        self.energy = -1
        self.gamma = gamma
        self.alpha = alpha
        self.last_choice_index = None
        self.weights = [1.0 for _ in range(len(self.fuzzer.templates))]
        self.probs = [0.0 for _ in range(len(self.fuzzer.templates))]

    def select(self) -> "Template":
        if len(self.fuzzer.templates) > len(self.weights):
            self.weights.extend(
                [1.0 for _ in range(len(self.fuzzer.templates) - len(self.weights))]
            )

        np_weights = np.array(self.weights)
        probs = (1 - self.gamma) * np_weights / np_weights.sum() + self.gamma / len(
            self.fuzzer.templates
        )

        self.last_choice_index = np.random.choice(len(self.fuzzer.templates), p=probs)

        self.fuzzer.templates[self.last_choice_index].visited_num += 1
        self.probs[self.last_choice_index] = probs[self.last_choice_index]

        return self.fuzzer.templates[self.last_choice_index]

    def update(self, template: "Template"):
        succ_num = template.num_jailbreaks

        r = 1 - succ_num
        x = -1 * r / self.probs[self.last_choice_index]
        self.weights[self.last_choice_index] *= np.exp(self.alpha * x / len(self.fuzzer.templates))


class QLearningSelectPolicy(SelectPolicy):
    def __init__(self, fuzzer: "TurboFuzzLLMFuzzer" = None):
        super().__init__(fuzzer)
        self.q_table: "dict[str, dict[Template, float]]" = dict()
        self.curr_action = None

        # Define hyper parameters
        self.learning_rate = 0.3
        self.discount_factor = 0.02
        self.epsilon = 1.0
        self.epsilon_decay_factor = 0.97
        self.epsilon_min = 0.2
        self.diversity_factor = 0.08

    def log_q_table(self):
        LOGGER.debug(f"\t== TEMPLATE 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 initialize(self):
        state = ""
        self.q_table[state] = dict()
        for action in self.fuzzer.initial_templates:
            if action not in self.q_table[state]:
                self.q_table[state][action] = action.initial_asr

    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)
        if self.fuzzer.mode == "attack":
            return (100.0 * num_success) / (num_success + num_failure)
        else:
            success_freq_discount = math.pow(
                1.0 - self.diversity_factor, gen_template.root_parent.visited_num
            )
            failure_freq_amplifier = math.pow(
                1.0 + self.diversity_factor, gen_template.root_parent.visited_num
            )
            return (
                ((100.0 * num_success) / (num_success + num_failure)) * success_freq_discount
            ) - (((100.0 * num_failure) / (num_success + num_failure)) * failure_freq_amplifier)

    def select(self) -> "Template":
        self.log_q_table()

        state = ""
        state_table = self.q_table.get(state)

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

        # decay epsilon
        self.epsilon = max(self.epsilon_min, self.epsilon_decay_factor * self.epsilon)

        self.curr_action.root_parent.visited_num += 1
        return self.curr_action

    def _update(self, state, action, reward):
        orig = self.q_table[state][action]
        new = (1 - self.learning_rate) * orig + self.learning_rate * (
            reward + self.discount_factor * max(self.q_table[state].values())
        )
        LOGGER.debug(
            f"\tTEMPLATE SELECTOR Q-LEARNING UPDATE: Reward {reward} changed q value for ({state}, {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].values())}"
            )
            LOGGER.debug(f"NEW: {new}")
        self.q_table[state][action] = new

    def update(self, gen_template: "Template"):
        state = ""

        # add new/generated template to Q table
        assert gen_template not in self.q_table[state]
        if QUESTION_PLACEHOLDER in gen_template.text:
            # if gen_template.num_jailbreaks > 0:
            self.q_table[state][gen_template] = gen_template.asr

        # reward selected/parent template
        reward = self.calculate_reward(gen_template)
        self._update(state, self.curr_action, reward)
