"""
random_suffix_search.py

This attack implements a random suffix search attack by operating at the token level.
It is tailored for OpenAI models by using tiktoken's "cl100k_base" encoding directly.
The attack works as follows:
  - An initial adversarial suffix of ~25 tokens is generated by sampling uniformly
    from the token vocabulary.
  - At each iteration, a random contiguous block (up to a maximum of 4 tokens) in the
    suffix is replaced with new random tokens.
  - The modified suffix is inserted into the original prompt:
        • If a "payload" field exists in the entry, the suffix is appended immediately
          after the payload.
        • Otherwise, the suffix is appended at the end of the document.
  - The resulting candidate is evaluated using the provided judge function.
This version does not use log probabilities to direct the search—it simply performs random
modifications and assesses success.

Usage:
    spikee test --attack random_suffix_search --attack-iterations 100

Parameters:
    entry (dict):
        The dataset entry containing at least "text" and optionally "system_message" and "payload".
        May also include "n_tokens_adv" (default 25) and "n_tokens_change_max" (default 4).
    target_module (module):
        The target module implementing process_input(input_text, system_message, logprobs=False).
    call_judge (function):
        A function to judge if the LLM response indicates success.
        Whether to treat the judge output as a boolean.
    max_iterations (int):
        Maximum number of attack iterations.
    attempts_bar (tqdm, optional):
        A progress bar to update with each iteration.

Returns:
    tuple: (iterations_attempted, success_flag, modified_input, last_response)
"""

import random
import time
import numpy as np
import tiktoken
import threading

# Default parameters.
n_tokens_adv = 25   # Number of tokens in the adversarial suffix.
n_tokens_change = 4 # Maximum tokens to change per iteration.

def insert_adv_string(document, adv_string, payload=None):
    """
    Inserts the adversarial string into the document.
    If a payload is provided and found in the document, the adversarial string is appended
    immediately after the first occurrence of the payload.
    Otherwise, the adversarial string is appended to the end of the document.
    """
    if payload:
        idx = document.find(payload)
        if idx != -1:
            return document[:idx + len(payload)] + " " + adv_string + document[idx + len(payload):]
    return document + "\n" + adv_string

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)
    payload_field = entry.get("payload", "")
    
    # Use tiktoken directly for OpenAI models.
    tokenizer = tiktoken.get_encoding("cl100k_base")
    max_token_value = tokenizer.n_vocab
    
    # Initialize the adversarial tokens as a random sequence.
    adv_tokens = np.random.randint(0, max_token_value, n_tokens_adv).tolist()
    
    # Build the initial candidate prompt.
    candidate = insert_adv_string(original_text, tokenizer.decode(adv_tokens), payload=payload_field)
    response = ""
    
    # Main random search loop.
    for i in range(1, max_iterations + 1):
        # Randomly choose a starting index for substitution.
        if n_tokens_adv >= n_tokens_change:
            substitute_pos_start = random.randint(0, n_tokens_adv - n_tokens_change)
        else:
            substitute_pos_start = 0
        
        # Generate new random tokens for substitution.
        substitution_tokens = np.random.randint(0, max_token_value, n_tokens_change).tolist()
        # Replace the selected tokens in adv_tokens.
        adv_tokens[substitute_pos_start:substitute_pos_start + n_tokens_change] = substitution_tokens
        
        # Decode the modified token sequence into text.
        adv = tokenizer.decode(adv_tokens).replace('<s>', '')
        # Insert the adversarial suffix into the document.
        candidate = insert_adv_string(original_text, adv, payload=payload_field)
            
        try:
            # Call process_input. The wrapper guarantees a tuple is returned.
            response, _ = target_module.process_input(candidate, system_message, logprobs=False)
            success = call_judge(entry, response)
        except Exception as e:
            success = False
            response = str(e)
        
        # Update progress bar safely.
        if attempts_bar:
            with bar_lock:
                attempts_bar.update(1)
                
        if success:
            if attempts_bar:
                with bar_lock:
                    remaining = max_iterations - i
                    attempts_bar.total = attempts_bar.total - remaining
            return i, True, candidate, response
    
    return max_iterations, False, candidate, response
