"""
best-of-n.py

This attack implements a Best-of-N random scramble attack based on a random scramble
augmentation strategy (e.g., from "Best-of-N Jailbreaking" by John Hughes et al., 2024).

Attack Interface:
    attack(entry, target_module, call_judge, max_iterations, attempts_bar=None)

Arguments:
    entry: dict
        The dataset item (must include "text" and preferably a "payload" field and
        "exclude_from_transformations_regex" which is a list of regex strings).
    target_module: module
        The target module providing process_input(input_text, system_message).
    call_judge: function
        A function that judges the LLM output: call_judge(entry, llm_output) -> bool.
    max_iterations: int
        Maximum number of attack iterations to try.
    attempts_bar: tqdm progress bar (optional)
        If provided, the attack will update it during its iterations.

Returns:
    A tuple: (iterations_attempted, success_flag, last_payload, last_response)
"""

import re
import random
import threading

# Transformation probability constants
SCRAMBLE_PROB = 0.6         
CAPITALIZATION_PROB = 0.6     
NOISE_PROB = 0.06           

def attack(entry, target_module, call_judge, max_iterations, attempts_bar=None, bar_lock=None):
    original_text = entry.get("text", "")
    system_message = entry.get("system_message", None)
    # Use the payload field (if present) to restrict transformation.
    payload_field = entry.get("payload", "")
    # Retrieve the exclusion list from the entry.
    exclusions = entry.get("exclude_from_transformations_regex", [])
    last_payload = original_text  # fallback if no transformation occurs
    last_response = ""
        
    for i in range(1, max_iterations + 1):
        if payload_field:
            candidate_text = _scramble_payload_only(original_text, payload_field, exclusions)
        else:
            candidate_text = _generate_variant(original_text, exclusions)
        last_payload = candidate_text

        try:
            response, _ = target_module.process_input(candidate_text, system_message)
            last_response = response
            success = call_judge(entry, response)
        except Exception as e:
            success = False
            last_response = str(e)
            print(f"[Best-Of-N] Entry ID {entry['id']}: {e}")

        if attempts_bar:
            with bar_lock:
                attempts_bar.update(1)
                
        if success:
            # If successful, update the progress bar.
            if attempts_bar:
                with bar_lock:
                    remaining = max_iterations - i
                    attempts_bar.total = attempts_bar.total - remaining
            return i, True, candidate_text, response

    return max_iterations, False, last_payload, last_response

def _scramble_payload_only(original_text, payload, exclusions):
    """
    Finds the first occurrence of `payload` in `original_text` and returns a new string
    where that occurrence is replaced with a scrambled variant generated by _generate_variant.
    If the payload is not found, falls back to scrambling the entire original_text.
    """
    idx = original_text.find(payload)
    if idx == -1:
        return _generate_variant(original_text, exclusions)
    scrambled_payload = _generate_variant(payload, exclusions)
    return original_text[:idx] + scrambled_payload + original_text[idx+len(payload):]

def _generate_variant(text, exclusions):
    """
    Generates a variant of the input text by applying a random scramble transformation.
    If exclusions is provided, it builds a compound regex from them and uses re.split()
    so that any substring that exactly matches one of these regexes is left unchanged.
    """
    if exclusions:
        compound = "(" + "|".join(exclusions) + ")"
        chunks = re.split(compound, text)
        compound_re = re.compile(compound)
    else:
        chunks = [text]
        compound_re = None

    result_chunks = []
    for chunk in chunks:
        if compound_re and compound_re.fullmatch(chunk):
            # This chunk exactly matches one of the exclusion patterns: leave it unchanged.
            result_chunks.append(chunk)
        else:
            result_chunks.append(_augment_text(chunk))
    return "".join(result_chunks)

def _augment_text(text):
    """
    Applies the following transformations:
      - For tokens of length ≥ 4, scramble the middle letters with probability SCRAMBLE_PROB.
      - Randomly capitalize letters with probability CAPITALIZATION_PROB.
      - Apply character noising (ASCII ±1) with probability NOISE_PROB.
    The text is split on whitespace to preserve spacing.
    """
    tokens = re.split(r'(\s+)', text)
    transformed_tokens = []
    for token in tokens:
        if token.strip() == "":
            transformed_tokens.append(token)
        else:
            transformed_tokens.append(_maybe_scramble_words(token))
    scrambled = "".join(transformed_tokens)
    result_chars = []
    for c in scrambled:
        if random.random() < CAPITALIZATION_PROB:
            c = c.upper()
        code = ord(c)
        if 32 <= code <= 126 and random.random() < NOISE_PROB:
            delta = random.choice([-1, 1])
            new_code = code + delta
            if 32 <= new_code <= 126:
                c = chr(new_code)
        result_chars.append(c)
    return "".join(result_chars)

def _maybe_scramble_words(token):
    """
    For each subword (split on non-alphanumeric characters) in the token,
    scramble the middle letters (if the subword length ≥ 4 and with probability SCRAMBLE_PROB).
    """
    subwords = re.split(r'([^a-zA-Z0-9]+)', token)
    scrambled = []
    for sub in subwords:
        if not sub or re.fullmatch(r'[^a-zA-Z0-9]+', sub):
            scrambled.append(sub)
        else:
            if len(sub) >= 4 and random.random() < SCRAMBLE_PROB:
                mid = list(sub[1:-1])
                random.shuffle(mid)
                sub = sub[0] + "".join(mid) + sub[-1]
            scrambled.append(sub)
    return "".join(scrambled)
