import csv
import random
from concurrent import futures
from concurrent.futures import ThreadPoolExecutor

from pandas import Series

from turbofuzzllm.fuzzer.question import Question
from turbofuzzllm.fuzzer.template import Response, Template
from turbofuzzllm.utils.draw import draw_template_tree
from turbofuzzllm.utils.logging import LOGGER
from turbofuzzllm.utils.prompt import QUESTION_PLACEHOLDER, synthesis_message
from turbofuzzllm.utils.statistics import *
from turbofuzzllm.utils.timer import get_elapsed_time


class TurboFuzzLLMFuzzer:
    def __init__(
        self,
        questions: "dict[int, Series]",
        attacker,
        target,
        judge,
        initial_templates: "dict[str, Series]",
        mutate_policy: "MutatePolicy",
        select_policy: "SelectPolicy",
        max_queries: int = -1,
        max_warmup_queries: int = -1,
        max_jailbreaks_per_question: int = -1,
        out_path: "str" = "test",
        mode: str = "attack",
        num_threads: int = 1,
        guardrails_enabled: bool = False,
        early_stopping_constant: int = 4,
    ):

        self.all_questions: "dict[int, Question]" = {}
        self.remaining_questions: "list[int]" = []

        for idx, row in questions.items():
            question = Question(text=row["text"], idx=idx, data=row)
            self.all_questions[idx] = question
            self.remaining_questions.append(idx)
        stats().write("#-questions", len(questions))

        self.jailbreaks: "dict[int, dict[str, set[str]]]" = {}

        self.attacker = attacker
        self.target = target
        self.judge = judge

        self.all_templates: "list[Template]" = []
        self.initial_templates: "list[Template]" = []
        self.templates: "list[Template]" = []
        for name, row in initial_templates.items():
            template = Template(
                text=row["text"],
                name=name,
                data=row,
                unexplored_child_mutations=mutate_policy.mutators.copy(),
            )
            self.all_templates.append(template)
            if QUESTION_PLACEHOLDER in template.text:
                self.initial_templates.append(template)
                self.templates.append(template)
        stats().write("#-templates-original", len(initial_templates))

        self.mutate_policy = mutate_policy
        self.select_policy = select_policy
        self.mutate_policy.fuzzer = self
        self.select_policy.fuzzer = self

        self.num_attacks: int = 0
        self.num_rejects: int = 0
        self.num_iterations: int = 0
        self.num_invalid_templates = 0

        self.max_queries: int = max_queries
        self.max_warmup_queries: int = max_warmup_queries
        self.max_jailbreaks_per_question: int = max_jailbreaks_per_question

        self.out_path = out_path
        self.mode: str = mode
        self.num_threads: int = num_threads

        self.guardrails_enabled = guardrails_enabled
        self.guardrails_safe_blocked = 0
        self.guardrails_vulnerable_blocked = 0
        self.guardrails_safe_unblocked = 0
        self.guardrails_vulnerable_unblocked = 0
        self.early_stopping_constant = early_stopping_constant

        self.writers = dict()

        fp = open(f"{self.out_path}/queries.csv", "w", buffering=1)
        self.writers["queries"] = (csv.writer(fp), fp)
        self.writers["queries"][0].writerow(["template", "qid", "question", "response", "result"])
        LOGGER.info(f"Recording queries in file queries.csv")

        fp = open(f"{self.out_path}/templates.csv", "w", buffering=1)
        self.writers["templates"] = (csv.writer(fp), fp)
        self.writers["templates"][0].writerow(
            ["template", "mutator", "parent", "root", "text", "#-attacks", "#-jailbreaks", "asr"]
        )
        LOGGER.info(f"Summarizing templates in file templates.csv")

        if mode != "evaluate":
            fp = open(f"{self.out_path}/mutators.csv", "w", buffering=1)
            self.writers["mutators"] = (csv.writer(fp), fp)
            self.writers["mutators"][0].writerow(
                [
                    "mutator",
                    "#-mutants",
                    "#-invalid-mutants",
                    "asr_num_jailbreaks",
                    "asr_num_jailbreaking_mutants",
                ]
            )
            LOGGER.info(f"Summarizing mutators in file mutators.csv")

    def is_stop(self):
        return (len(self.remaining_questions) == 0) or (
            self.max_queries != -1 and stats().num_queries >= self.max_queries
        )

    def write_mutators(self):
        for mutator in self.mutate_policy.mutators:
            self.writers["mutators"][0].writerow(
                [
                    mutator,
                    mutator.num_mutants,
                    mutator.num_invalid_mutants,
                    mutator.asr_num_jailbreaks,
                    mutator.asr_num_jailbreaking_mutants,
                ]
            )

    def write_stats(self):
        stats().write("time-seconds", get_elapsed_time())

        stats().write("seed-asr-top-1", self.get_seed_asr_top_n(1))
        stats().write("seed-asr-top-5", self.get_seed_asr_top_n(5))
        stats().write("seed-asr-top-10", self.get_seed_asr_top_n(10))
        stats().write("seed-asr-top-20", self.get_seed_asr_top_n(20))
        stats().write("seed-asr-all", self.seed_asr)

        stats().write("unseeded-asr-top-1", self.get_unseed_asr_top_n(1))
        stats().write("unseeded-asr-top-5", self.get_unseed_asr_top_n(5))
        stats().write("unseeded-asr-top-10", self.get_unseed_asr_top_n(10))
        stats().write("unseeded-asr-top-20", self.get_unseed_asr_top_n(20))
        stats().write("unseeded-asr-all", self.unseed_asr)

        stats().write("#-iterations", self.num_iterations)
        stats().write("#-attacks", self.num_attacks)
        stats().write("#-queries", stats().num_queries)
        stats().write("#-jailbreaks", self.get_total_jailbreaks())
        stats().write("#-rejects", self.num_rejects)
        stats().write("#-questions-jailbroken", len(self.jailbreaks))

        stats().write("#-templates-final", len(self.templates))
        stats().write("#-templates-all", len(self.all_templates))
        stats().write("#-templates-successful", self.get_num_successful_templates())
        stats().write("#-templates-invalid", self.num_invalid_templates)

        stats().write("guardrails-#-blocked", self.num_guardrails_blocked)
        stats().write("guardrails-#-unblocked", self.num_guardrails_unblocked)
        stats().write("guardrails-%-vulnerable-blocked", self.percent_vulnerable_blocks)
        stats().write("guardrails-%-vulnerable-unblocked", self.percent_vulnerable_unblocks)

    def close_writers(self):
        stats().close()
        for _, fp in self.writers.values():
            fp.close()

    def close(self):
        self.log(True, True)
        draw_template_tree(
            self.all_templates,
            self.all_questions,
            f"{self.out_path}/template_tree.dot",
            False,
        )
        if self.mode != "evaluate":
            self.write_mutators()
        self.write_stats()
        self.attacker.cleanup()
        self.target.cleanup()
        self.judge.cleanup()
        self.close_writers()

    @property
    def seed_asr(self):
        return (100.0 * len(self.jailbreaks)) / len(self.all_questions)

    @property
    def unseed_asr(self):
        if self.num_attacks == 0:
            return 0.0
        return (100.0 * self.get_total_jailbreaks()) / self.num_attacks

    @property
    def num_guardrails_blocked(self):
        return self.guardrails_vulnerable_blocked + self.guardrails_safe_blocked

    @property
    def num_guardrails_unblocked(self):
        return self.guardrails_vulnerable_unblocked + self.guardrails_safe_unblocked

    @property
    def percent_vulnerable_blocks(self):
        if self.num_guardrails_blocked == 0:
            return 0.0
        return (100.0 * self.guardrails_vulnerable_blocked) / self.num_guardrails_blocked

    @property
    def percent_vulnerable_unblocks(self):
        if self.num_guardrails_unblocked == 0:
            return 0.0
        return (100.0 * self.guardrails_vulnerable_unblocked) / self.num_guardrails_unblocked

    def get_total_jailbreaks_for_question(self, qid):
        if qid not in self.jailbreaks:
            return 0
        result = 0
        for _, v in self.jailbreaks.get(qid).items():
            result += len(v)
        return result

    def get_total_jailbreaks(self):
        result = 0
        for qid in self.all_questions.keys():
            result += self.get_total_jailbreaks_for_question(qid)
        return result

    def get_num_successful_templates(self):
        top_n = [t for t in self.all_templates if t.num_attacks > 0 and t.num_jailbreaks > 0]
        return len(top_n)

    def print_top_n_templates(self, n):
        top_n = [t for t in self.all_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]
        return "".join([f"{x} -> {x.num_jailbreaks}/{x.asr:.1f}%/{x.root}, " for x in top_n])

    def print_top_n_mutations(self, n):
        top_n = [m for m in self.mutate_policy.mutators if m.num_mutants > 0]
        top_n = sorted(
            top_n,
            key=lambda x: (x.asr_num_jailbreaks, x.asr_num_jailbreaking_mutants, x.num_mutants),
            reverse=True,
        )
        if n > 0:
            top_n = top_n[:n]
        return "".join(
            [
                f"{x} -> {x.asr_num_jailbreaks:.1f}%/{x.asr_num_jailbreaking_mutants:.1f}%, "
                for x in top_n
            ]
        )

    def get_seed_asr_top_n(self, n):
        top_n = [t for t in self.all_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]
        jailbroken_questions = set()
        for pn in top_n:
            for qid in pn.jailbroken_questions:
                jailbroken_questions.add(qid)
        return (100.0 * len(jailbroken_questions)) / len(self.all_questions)

    def get_unseed_asr_top_n(self, n):
        top_n = [t for t in self.all_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]
        nattacks = 0
        njailbreaks = 0
        for pn in top_n:
            nattacks += pn.num_attacks
            njailbreaks += pn.num_jailbreaks
        if nattacks == 0:
            return 0.0
        return (100.0 * njailbreaks) / nattacks

    def run(self):
        if self.mode == "evaluate":
            self.run_input_eval()
        elif self.mode == "legacy":
            self.run_warmup()
            self.run_mutation()
        elif self.mode == "attack" or self.mode == "answer":
            self.run_warmup()
            self.run_mutation()
        else:
            raise Exception(f"Unexpected mode: {self.mode}")

    def run_input_eval(self):
        LOGGER.info(f"Starting {self.mode} mode!")
        try:
            for original in self.all_templates:
                self.num_iterations += 1
                LOGGER.info(f"{self.num_iterations}. TEMPLATE: {original}")
                LOGGER.debug(f"== {original} ==\n{original.text}\n==")
                if QUESTION_PLACEHOLDER not in original.text:
                    LOGGER.info(f"\tORIGINAL TEMPLATE: INVALID")
                    self.num_invalid_templates += 1
                    continue
                self.evaluate(original, self.remaining_questions)
                self.update(original)
                self.log(False, False)
                if self.is_stop():
                    break
        except KeyboardInterrupt:
            LOGGER.info("TurboFuzzLLM interrupted by user!")
        LOGGER.info(f"  Finished {self.mode} mode!")
        self.close()

    def run_warmup(self):
        if self.max_warmup_queries < len(self.all_templates):
            return
        num_warmup_questions = int(self.max_warmup_queries / len(self.all_templates))
        if self.mode == "attack":
            num_warmup_questions = len(self.remaining_questions)

        assert num_warmup_questions > 0
        LOGGER.info(f"Starting warmup in {self.mode} mode using {num_warmup_questions} questions")

        try:
            for original in sorted(self.all_templates, key=lambda x: x.initial_asr, reverse=True):
                self.num_iterations += 1
                LOGGER.info(f"{self.num_iterations}. TEMPLATE: {original}")
                LOGGER.debug(f"== {original} ==\n{original.text}\n==")
                if QUESTION_PLACEHOLDER not in original.text:
                    LOGGER.info(f"\tORIGINAL TEMPLATE: INVALID")
                    self.num_invalid_templates += 1
                    continue
                # select at most `num_warmup_questions` out of remaining questions randomly
                if len(self.remaining_questions) > num_warmup_questions:
                    warmup_questions = sorted(
                        self.remaining_questions, key=lambda x: random.random()
                    )[:num_warmup_questions]
                else:
                    warmup_questions = self.remaining_questions
                self.evaluate(original, warmup_questions)
                self.update(original)
                original.initial_asr = original.asr
                self.log(False, False)
                if self.is_stop() or stats().num_queries > self.max_warmup_queries:
                    break
        except KeyboardInterrupt:
            LOGGER.info("Warmup interrupted by user!")
        LOGGER.info("Warmup finished!")

    def run_mutation(self):
        LOGGER.info(f"Starting fuzzing for {self.mode} mode!")

        self.select_policy.initialize()
        try:
            while not self.is_stop():
                self.num_iterations += 1
                original = self.select_policy.select()
                LOGGER.info(f"{self.num_iterations}. TEMPLATE: {original}")
                LOGGER.debug(f"== {original} ==\n{original.text}\n==")
                mutator, status = self.mutate_policy.select(original)
                if status != "Succeeded":
                    continue
                LOGGER.info(f"\tMUTATOR: {mutator}")
                mutant = self.mutate_policy.mutate_single(original, mutator)
                LOGGER.debug(f"== {mutant} ==\n{mutant.text}\n==")
                if QUESTION_PLACEHOLDER not in mutant.text:
                    LOGGER.info(f"\tMUTATED TEMPLATE: INVALID")
                    self.num_invalid_templates += 1
                    self.select_policy.update(mutant)
                    self.mutate_policy.update(mutant)
                    continue
                self.all_templates.append(mutant)
                LOGGER.info(f"\tMUTATED TEMPLATE: {mutant}/{mutant.root}")
                self.evaluate(mutant, self.remaining_questions)
                self.select_policy.update(mutant)
                self.mutate_policy.update(mutant)
                self.update(mutant)
                self.log(False, True)
        except KeyboardInterrupt:
            LOGGER.info("TurboFuzzLLM interrupted by user!")
        LOGGER.info(f"  Finished {self.mode} mode!")
        self.close()

    def evaluate(self, template: "Template", eval_questions_orig: "list[int]"):
        LOGGER.info(f"\tEVALUATING: {template}")
        eval_questions = eval_questions_orig
        num_questions_per_batch = len(eval_questions)
        if self.mode == "attack" or self.mode == "answer":
            # we just want to break each question once
            # batch with 10% random questions, with minimum 4 and maximum 20
            num_questions_per_batch = max(
                self.early_stopping_constant, min(20, int(len(eval_questions) * 0.1))
            )
            # select a random subset of eval questions
            eval_questions = sorted(eval_questions, key=lambda x: random.random())
        i = 0
        while i < len(eval_questions):
            eval_questions_in_batch = eval_questions[i : i + num_questions_per_batch]
            if self.num_threads > 1:
                num_jb = self.evaluate_multi_process(template, eval_questions_in_batch)
            else:
                num_jb = self.evaluate_single_process(template, eval_questions_in_batch)
            i = i + num_questions_per_batch
            if (self.mode == "attack" or self.mode == "answer") and num_jb == 0:
                break
        if self.mode != "answer":
            LOGGER.info(f"\t{template} ASR: {template.asr:.1f}%/{template.root}")

    def evaluate_single_process(self, template: "Template", eval_questions: "list[int]"):
        num_jb = 0
        for qid in eval_questions:
            if qid not in template.responses:
                num_jb += self.evaluate_question_single(template, qid)
        return num_jb

    def evaluate_multi_process(self, template: "Template", eval_questions: "list[int]"):
        num_jb = 0
        for idx, njb, exception in self.evaluate_questions_multi(template, eval_questions):
            qid = eval_questions[idx]
            if exception is None:
                num_jb += njb
            else:
                assert (
                    False
                ), f"Bedrock API call failed due to {exception} for question {self.all_questions[qid]}."
        return num_jb

    def evaluate_question_single(self, template: "Template", qid: int):
        question = self.all_questions[qid]
        LOGGER.debug(f"== QUESTION {question} ==\n{question.text}\n==")
        message = synthesis_message(question.text, template.text)
        assert message is not None
        stats().num_queries += 1
        response = self.target.generate(message, f"{question}")
        num_jb = self.evaluate_question_response(template, qid, response)
        return num_jb

    def evaluate_questions_multi(self, template: "Template", eval_questions: "list[int]"):
        with ThreadPoolExecutor(max_workers=self.num_threads) as executor:
            future_to_idx = {
                executor.submit(self.evaluate_question_single, template, eval_questions[idx]): idx
                for idx in range(len(eval_questions))
            }

            for future in futures.as_completed(future_to_idx):
                idx = future_to_idx[future]
                exception = future.exception()

                if not exception:
                    yield idx, future.result(), None
                else:
                    yield idx, None, exception

    def evaluate_question_response(
        self, template: "Template", qid: int, response: "dict[str, str]"
    ):
        question = self.all_questions[qid]
        result = self.judge.judge(question, response["text"])
        if self.guardrails_enabled:
            if response["status"] == "guardrail_intervened":
                result = "GUARDRAIL_BLOCKED_" + result
                if result == "VULNERABLE":
                    self.guardrails_vulnerable_blocked += 1
                else:
                    self.guardrails_safe_blocked += 1
            else:
                if result == "VULNERABLE":
                    self.guardrails_vulnerable_unblocked += 1
                else:
                    self.guardrails_safe_unblocked += 1
        LOGGER.info(f"\t\t{question}: {result}")
        assert qid not in template.responses
        template.responses[qid] = Response(text=response["text"], result=result)
        if result == "VULNERABLE":
            template.jailbroken_questions.add(qid)
            if self.mode == "answer":
                print(f"== ORIGINAL QUESTION == {question.text}\n")
                templated_question = synthesis_message(question.text, template.text)
                print(f"======= PROMPT ======== {templated_question}\n")
                print(f"======= ANSWER ======== {response['text']}\n")
            return 1
        return 0

    def update(self, template):
        if (template.num_jailbreaks > 0) and (template.parent != template):
            if QUESTION_PLACEHOLDER in template.text:
                self.templates.append(template)
        for qid, response in template.responses.items():
            question_text = self.all_questions[qid].text

            self.writers["queries"][0].writerow(
                [template, qid, question_text, response.text, response.result]
            )

            if response.result == "VULNERABLE":
                if qid not in self.jailbreaks:
                    self.jailbreaks[qid] = dict()
                if template.root not in self.jailbreaks[qid]:
                    self.jailbreaks[qid][template.root] = set()
                self.jailbreaks[qid][template.root].add(template.name)
                if (
                    0
                    < self.max_jailbreaks_per_question
                    <= self.get_total_jailbreaks_for_question(qid)
                ):
                    if qid in self.remaining_questions:
                        self.remaining_questions.remove(qid)

        self.num_attacks += template.num_attacks
        self.num_rejects += template.num_rejects

        self.writers["templates"][0].writerow(
            [
                template,
                template.mutator,
                template.parent,
                template.root,
                template.text,
                template.num_attacks,
                template.num_jailbreaks,
                template.asr,
            ]
        )

    def log(self, end_print: bool, print_mut: bool):
        if end_print or self.mode != "answer":
            LOGGER.info(
                f"  Summary after {get_elapsed_time():.0f} seconds and {stats().num_queries} queries: total {self.get_total_jailbreaks()} jailbreaks over {len(self.jailbreaks)} questions after {self.num_attacks} attacks using {self.num_iterations} templates ({self.get_num_successful_templates()} successful, {self.num_invalid_templates} invalid)"
            )
            LOGGER.info(
                f"\tSeed ASR:        top-1 -> {self.get_seed_asr_top_n(1):.1f}%, top-5 -> {self.get_seed_asr_top_n(5):.1f}%, top-10 -> {self.get_seed_asr_top_n(10):.1f}%, all -> {self.seed_asr:.1f}%"
            )
            LOGGER.info(
                f"\tUnseeded ASR:    top-1 -> {self.get_unseed_asr_top_n(1):.1f}%, top-5 -> {self.get_unseed_asr_top_n(5):.1f}%, top-10 -> {self.get_unseed_asr_top_n(10):.1f}%, all -> {self.unseed_asr:.1f}%"
            )
            LOGGER.info(f"\tTop-5 Templates: {self.print_top_n_templates(5)}")
            if print_mut and (self.mode != "evaluate"):
                LOGGER.info(f"\tTop-3 Mutations: {self.print_top_n_mutations(3)}")
            if self.guardrails_enabled:
                LOGGER.info(
                    f"\tGuardrails:          {self.num_guardrails_blocked} blocked ({self.percent_vulnerable_blocks:.1f}% "
                    f"vulnerable), {self.num_guardrails_unblocked} unblocked ({self.percent_vulnerable_unblocks:.1f}% "
                    f"vulnerable)"
                )
