# Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Requirement: Any integration or derivative work must explicitly attribute
# Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
# documentation or user interface, as detailed in the NOTICE file.

from tqdm.asyncio import tqdm as async_tqdm_bar
from pydantic import BaseModel
from tqdm import tqdm
import asyncio
import time
import json


from deepeval.models import DeepEvalBaseLLM

from deepteam.attacks import BaseAttack
from deepteam.attacks.multi_turn.tree_jailbreaking.schema import (
    ImprovementPrompt,
    NonRefusal,
    Rating,
)
from deepteam.attacks.multi_turn.tree_jailbreaking.template import (
    JailBreakingTemplate,
)
from deepteam.attacks.attack_simulator.utils import (
    generate_schema,
    a_generate_schema,
)
from deepteam.attacks.multi_turn.types import CallbackType


class TreeNode:
    def __init__(
        self, prompt: str, score: int, depth: int, conversation_history=None
    ):
        self.prompt = prompt
        self.score = score
        self.depth = depth
        self.children = []
        self.conversation_history = conversation_history or []


class TreeJailbreaking(BaseAttack):

    def __init__(
        self,
        weight: int = 1,
    ):
        self.weight = weight

    def enhance(
        self,
        attack: str,
        model_callback: CallbackType,
        simulator_model: DeepEvalBaseLLM,
    ) -> str:
        self.simulator_model = simulator_model
        self.model_callback = model_callback

        MAX_RUNTIME = 50.0  # Time limit in seconds
        root = TreeNode(prompt=attack, score=0, depth=0)
        start_time = time.time()

        # Create a sync progress bar
        pbar = tqdm(
            total=MAX_RUNTIME,
            desc="...... ⛓️  Tree Jailbreaking",
            unit="s",
            leave=False,
        )

        try:
            best_node = self.tree_search(
                root, attack, start_time, MAX_RUNTIME, pbar
            )
        finally:
            pbar.close()  # Make sure the progress bar is properly closed

        return best_node.prompt

    def tree_search(
        self,
        root: TreeNode,
        goal: str,
        start_time: float,
        MAX_RUNTIME: float,
        pbar,
    ) -> TreeNode:
        """Synchronous tree search implementation with runtime limit."""
        queue = [root]
        best_node = root

        while queue and (time.time() - start_time < MAX_RUNTIME):
            node = queue.pop(0)
            self.expand_node(node, goal)
            queue.extend(node.children)
            queue.sort(key=lambda x: x.score, reverse=True)

            if node.score > best_node.score:
                best_node = node

            # Update progress bar
            elapsed_time = time.time() - start_time
            pbar.n = elapsed_time
            pbar.refresh()

        return best_node

    def expand_node(
        self, node: TreeNode, initial_attack: str, PRUNE_THRESHOLD: int = 0
    ):
        """Expands a tree node by generating multiple children nodes based on branching logic."""
        if node.depth == 0:
            conversation_json = [
                {
                    "role": "system",
                    "content": JailBreakingTemplate.enhance_attack(
                        initial_attack
                    ),
                }
            ]
        else:
            conversation_json = node.conversation_history

        conversation_json.append(
            {
                "role": "user",
                "content": f"Current prompt: {node.prompt}\n\nOBJECTIVE: {initial_attack}\n\nSCORE: {node.score}",
            }
        )

        # Generate multiple branches
        num_branches = self.calculate_branches(node.score, node.depth)
        for _ in range(num_branches):
            try:
                conversation = json.dumps(conversation_json)
                res: ImprovementPrompt = self._generate_schema(
                    conversation, ImprovementPrompt
                )
                enhanced_attack = res.prompt

                # Check if enhanced attack is a refusal statement
                non_refusal_prompt = JailBreakingTemplate.non_refusal(
                    initial_attack, enhanced_attack
                )
                res: NonRefusal = self._generate_schema(
                    non_refusal_prompt, NonRefusal
                )
                classification = res.classification
                if classification == "Refusal":
                    continue  # Skip this child

                # Generate a response from the target LLM
                target_response = self.model_callback(enhanced_attack)

                # Calculate the score for the enhanced attack
                judge_prompt = JailBreakingTemplate.linear_judge(
                    initial_attack, enhanced_attack, target_response
                )
                res: Rating = self._generate_schema(judge_prompt, Rating)
                score = res.rating

                # Prune if the score is too low
                if score <= PRUNE_THRESHOLD:
                    continue  # Skip creating this child

                # Create a new child node
                child_node = TreeNode(
                    prompt=enhanced_attack,
                    score=score,
                    depth=node.depth + 1,
                    conversation_history=conversation_json,
                )
                node.children.append(child_node)
            except Exception as e:
                continue

    async def a_enhance(
        self,
        attack: str,
        model_callback: CallbackType,
        simulator_model: DeepEvalBaseLLM,
    ) -> str:
        self.model_callback = model_callback
        self.simulator_model = simulator_model

        MAX_RUNTIME = 50.0  # Time limit in seconds
        root = TreeNode(prompt=attack, score=0, depth=0)
        start_time = time.time()

        # Async progress bar task
        pbar = async_tqdm_bar(
            total=MAX_RUNTIME,
            desc="...... ⛓️  Tree Jailbreaking",
            unit="s",
            leave=False,
        )
        tick_task = asyncio.create_task(
            self.update_pbar(pbar, start_time, MAX_RUNTIME)
        )

        try:
            # Run tree search concurrently with the progress bar updates
            best_node = await self.a_tree_search(
                root, attack, start_time, MAX_RUNTIME
            )
        finally:
            # Make sure the progress bar is properly closed
            await tick_task

        return best_node.prompt

    async def a_tree_search(
        self, root: TreeNode, goal: str, start_time: float, MAX_RUNTIME: float
    ) -> TreeNode:
        """Asynchronous tree search implementation with runtime limit."""
        queue = [root]
        best_node = root

        while queue and (time.time() - start_time < MAX_RUNTIME):
            node = queue.pop(0)
            await self.a_expand_node(node, goal)
            queue.extend(node.children)
            queue.sort(key=lambda x: x.score, reverse=True)

            if node.score > best_node.score:
                best_node = node

        return best_node

    async def a_expand_node(
        self, node: TreeNode, initial_attack: str, PRUNE_THRESHOLD: int = 0
    ):
        """Expands a tree node asynchronously by generating multiple children nodes."""
        if node.depth == 0:
            conversation_json = [
                {
                    "role": "system",
                    "content": JailBreakingTemplate.enhance_attack(
                        initial_attack
                    ),
                }
            ]
        else:
            conversation_json = node.conversation_history

        conversation_json.append(
            {
                "role": "user",
                "content": f"Current prompt: {node.prompt}\n\nOBJECTIVE: {initial_attack}\n\nSCORE: {node.score}",
            }
        )

        # Generate branches asynchronously
        num_branches = self.calculate_branches(node.score, node.depth)
        tasks = [
            self.a_generate_child(
                node, conversation_json, initial_attack, PRUNE_THRESHOLD
            )
            for _ in range(num_branches)
        ]
        children = await asyncio.gather(*tasks, return_exceptions=True)

        # Filter out None values for pruned branches
        node.children.extend([child for child in children if (child and not isinstance(child, Exception))])

    async def a_generate_child(
        self,
        node: TreeNode,
        conversation_json: dict,
        initial_attack: str,
        PRUNE_THRESHOLD: int,
    ):
        """Asynchronously generates a child node."""
        conversation = json.dumps(conversation_json)
        res: ImprovementPrompt = await self._a_generate_schema(
            conversation, ImprovementPrompt
        )
        enhanced_attack = res.prompt

        # Check if enhanced attack is a refusal statement
        non_refusal_prompt = JailBreakingTemplate.non_refusal(
            initial_attack, enhanced_attack
        )
        res: NonRefusal = await self._a_generate_schema(
            non_refusal_prompt, NonRefusal
        )
        classification = res.classification
        if classification == "Refusal":
            return None  # Skip this child

        # Generate a response from the target LLM asynchronously
        target_response = await self.model_callback(enhanced_attack)

        # Calculate the score for the enhanced attack asynchronously
        judge_prompt = JailBreakingTemplate.linear_judge(
            initial_attack, enhanced_attack, target_response
        )
        res: Rating = await self._a_generate_schema(judge_prompt, Rating)
        score = res.rating

        # Prune if the score is too low
        if score <= PRUNE_THRESHOLD:
            return None

        # Create a new child node
        return TreeNode(
            prompt=enhanced_attack,
            score=score,
            depth=node.depth + 1,
            conversation_history=conversation_json,
        )

    ##################################################
    ### General Jailbreaking #########################
    ##################################################

    async def update_pbar(
        self, pbar: async_tqdm_bar, start_time: float, MAX_RUNTIME: float
    ):
        """Update the async progress bar every second and ensure it completes."""
        while True:
            elapsed_time = time.time() - start_time

            # Stop updating if the maximum runtime is exceeded
            if elapsed_time >= MAX_RUNTIME:
                pbar.n = MAX_RUNTIME
                pbar.refresh()
                pbar.close()
                break

            pbar.n = elapsed_time
            pbar.refresh()
            await asyncio.sleep(1)

    def calculate_branches(self, score: int, depth: int) -> int:
        """Calculate the number of branches (children) based on the score."""
        BASE_BRANCHES = 3
        MAX_BRANCHES = 5
        MIN_BRANCHES = 1

        branches = BASE_BRANCHES

        # Increase branches for high scores
        if score >= 8:
            branches += 2
        elif score >= 6:
            branches += 1

        # Decrease branches for low scores
        if score <= 3:
            branches -= 1

        return max(MIN_BRANCHES, min(MAX_BRANCHES, branches))

    ##################################################
    ### Utils ########################################
    ##################################################

    def _generate_schema(self, prompt: str, schema: BaseModel):
        return generate_schema(prompt, schema, self.simulator_model)

    async def _a_generate_schema(self, prompt: str, schema: BaseModel):
        return await a_generate_schema(prompt, schema, self.simulator_model)

    def get_name(self) -> str:
        return "Tree Jailbreaking"