import os
import re
import json
import time
import importlib
import random
import threading
import inspect
import traceback
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

class AdvancedTargetWrapper:
    """
    A wrapper for a target module's process_input method that incorporates both:
      - A loop for a given number of independent attempts (num_attempts), and
      - A retry strategy for handling 429 errors (max_retries) with throttling.

    This is designed to be passed to the attack() function so that each call to process_input()
    will try up to num_attempts times (each with up to max_retries on quota errors) before failing.

    Parameters:
      target_module: The original target module that provides process_input(input_text, system_message[, logprobs]).
      target_options: Target options, typically a string representing the name of the llm to call
      num_attempts (int): Number of independent attempts to call process_input per invocation.
      max_retries (int): Maximum number of retries per attempt (e.g. on 429 errors).
      throttle (float): Number of seconds to wait after a successful call.
    """
    def __init__(self, target_module, target_options=None, max_retries=3, throttle=0):
        self.target_module = target_module
        self.target_options = target_options
        self.max_retries = max_retries
        self.throttle = throttle

        sig = inspect.signature(self.target_module.process_input)
        params = sig.parameters
        # detect optional parameters that were only added in newer Spikee versions
        self.supports_options = 'target_options' in params
        self.supports_logprobs = 'logprobs' in params

    def process_input(self, input_text, system_message=None, logprobs=False):
        last_error = None
        retries = 0

        while retries < self.max_retries:
            try:
                # Build only the kwargs the underlying target supports.
                # Older targets without these parameters will simply be called without them.
                kwargs = {}
                if self.supports_options and self.target_options is not None:
                    kwargs['target_options'] = self.target_options
                if self.supports_logprobs:
                    kwargs['logprobs'] = logprobs

                # Delegate to the wrapped process_input
                if kwargs:
                    result = self.target_module.process_input(
                        input_text, system_message, **kwargs
                    )
                else:
                    result = self.target_module.process_input(
                        input_text, system_message
                    )

                # Unpack (response, logprobs) if tuple returned
                if isinstance(result, tuple) and len(result) == 2:
                    response, lp = result
                else:
                    response, lp = result, None

                if self.throttle > 0:
                    time.sleep(self.throttle)

                return response, lp

            except Exception as e:
                last_error = e
                if "429" in str(e) and retries < self.max_retries - 1:
                    wait_time = random.randint(30, 120)
                    time.sleep(wait_time)
                    retries += 1
                else:
                    break

        # All retries exhausted
        raise last_error

def validate_tag(tag):
    """
    Validates that a tag is safe to use in a filename.
    
    Args:
        tag (str): The tag to validate
        
    Returns:
        tuple: (is_valid, error_message)
            - is_valid (bool): True if tag is valid, False otherwise
            - error_message (str): Reason for validation failure or None if valid
    """
    if tag is None:
        return True, None
        
    # Check for empty string after stripping whitespace
    if len(tag.strip()) == 0:
        return False, "Tag cannot be empty or whitespace only"
    
    # Check length (reasonable max length for filename component)
    MAX_LENGTH = 50
    if len(tag) > MAX_LENGTH:
        return False, f"Tag exceeds maximum length of {MAX_LENGTH} characters"
    
    # Check for valid characters - alphanumeric, dash and underscore only
    pattern = re.compile(r'^[a-zA-Z0-9_-]+$')
    if not pattern.match(tag):
        return False, "Tag can only contain letters, numbers, dash (-) and underscore (_)"
    
    return True, None

def extract_dataset_name(file_name):
    file_name = os.path.basename(file_name)
    file_name = re.sub(r'^\d+-', '', file_name)
    file_name = re.sub(r'.jsonl$', '', file_name)
    if file_name.startswith("seeds-"):
        file_name = file_name[len("seeds-"):]
    return file_name

def read_jsonl_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        return [json.loads(line) for line in f]

def write_jsonl_file(output_file, data):
    with open(output_file, 'w', encoding='utf-8') as f:
        for entry in data:
            json.dump(entry, f, ensure_ascii=False)
            f.write('\n')

def append_jsonl_entry(output_file, entry, file_lock):
    """Appends a single entry to a JSONL file in a thread-safe manner."""
    with file_lock:
        with open(output_file, 'a', encoding='utf-8') as f:
            json.dump(entry, f, ensure_ascii=False)
            f.write('\n')

def check_canary_in_response(response, canary):
    if not canary:
        return False
    return canary in response

def load_module_from_path(path, module_name):
    spec = importlib.util.spec_from_file_location(module_name, path)
    if spec and spec.loader:
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        return mod
    raise ImportError(f"Could not load module {module_name} from {path}")

def _load_raw_target_module(target_name):
    """Load target module without wrapping. Returns None if not found."""
    local_path = os.path.join(os.getcwd(), 'targets', f"{target_name}.py")
    if os.path.isfile(local_path):
        return load_module_from_path(local_path, target_name)
    else:
        try:
            return importlib.import_module(f"spikee.targets.{target_name}")
        except ModuleNotFoundError:
            return None

def load_target_module(target_name, target_options, max_retries, throttle):
    target_mod = _load_raw_target_module(target_name)
    if target_mod is None:
        raise ValueError(f"Target '{target_name}' not found locally or in spikee.targets/")
    
    # Wrap the target module with AdvancedTargetWrapper
    return AdvancedTargetWrapper(target_mod, max_retries=max_retries, throttle=throttle, target_options=target_options)

def load_attack_by_name(attack_name):
    """
    Loads an attack module from a new "attacks" folder or from built-in package data.
    """
    local_attack_path = Path(os.getcwd()) / "attacks" / f"{attack_name}.py"
    if local_attack_path.is_file():
        spec = importlib.util.spec_from_file_location(attack_name, local_attack_path)
        attack_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(attack_module)
        return attack_module
    try:
        return importlib.import_module(f"spikee.attacks.{attack_name}")
    except ModuleNotFoundError:
        raise ValueError(f"Attack '{attack_name}' not found locally or in spikee.attacks")

def _load_raw_judge_module(judge_name):
    """Load judge module. Returns None if not found."""
    local_path = Path(os.getcwd()) / "judges" / f"{judge_name}.py"
    if local_path.is_file():
        spec = importlib.util.spec_from_file_location(judge_name, local_path)
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        return mod
    else:
        try:
            return importlib.import_module(f"spikee.judges.{judge_name}")
        except ModuleNotFoundError:
            return None

def _get_effective_attack_options(attack_module, provided_options):
    """Get effective attack options, using default if none provided and attack supports options."""
    if provided_options is not None:
        return provided_options
    
    # Try to get default option if none provided
    if attack_module and hasattr(attack_module, 'get_available_option_values'):
        try:
            available = attack_module.get_available_option_values()
            if available:
                return available[0]  # First option is default
        except Exception:
            pass
    
    return None

def _get_effective_judge_options(judge_name, provided_options):
    """Get effective judge options, using default if none provided and judge supports options."""
    if provided_options is not None:
        return provided_options
    
    # Try to get default option if none provided
    try:
        mod = _load_raw_judge_module(judge_name)
        if mod and hasattr(mod, 'get_available_option_values'):
            available = mod.get_available_option_values()
            if available:
                return available[0]  # First option is default
    except Exception:
        pass
    
    return None

def load_judge_module(judge_name):
    """
    Looks for `judges/{judge_name}.py` locally first,
    then falls back to built-in judges.
    """
    mod = _load_raw_judge_module(judge_name)
    if mod is None:
        raise ValueError(f"Judge '{judge_name}' not found locally or built-in.")
    return mod

def call_judge(entry, output):
    """
    Determines if the LLM output indicates a successful attack.
    
    If the output provided is a boolean that value is used to indicate success or failure. 
    This is used when testing LLM guardrail targets, which return True if the attack went 
    through the guardrail (attack successful) and False if the guardrail stopped it.
    
    In all other cases (i.e. when using a target LLM), the appropriate judge module
    for the attack is loaded and its judge() function is called.
    """
    if isinstance(output, bool):
        return output
      
    else:
        judge_name = entry.get("judge_name", "canary")
        judge_args = entry.get("judge_args", "")
        judge_options = entry.get("judge_options", None)
        llm_input = entry["text"]
        judge_module = load_judge_module(judge_name)
        return judge_module.judge(llm_input=llm_input, llm_output=output, judge_args=judge_args, judge_options=judge_options)

def _do_single_request(entry, input_text, target_module, num_attempt,
                        attempts_bar, global_lock):
    """
    Executes one request against the target by calling its process_input() method.
    The target_module is assumed to be an instance of AdvancedTargetWrapper that
    already implements retries and throttling. 

    Parameters:
      entry (dict): The dataset entry.
      input_text (str): The prompt text.
      target_module: The wrapped target module.
      num_attempt: The current attempt number.
      attempts_bar (tqdm): Progress bar to update.
      global_lock (threading.Lock): Lock for safely updating the progress bar.

    Returns:
      tuple: (result_dict, success)
    """
    # Extract metadata from the entry.
    task_type = entry.get("task_type", None)
    jailbreak_type = entry.get("jailbreak_type", None)
    instruction_type = entry.get("instruction_type", None)
    document_id = entry.get("document_id", None)
    position = entry.get("position", None)
    spotlighting_data_markers = entry.get("spotlighting_data_markers", None)
    injection_delimiters = entry.get("injection_delimiters", None)
    suffix_id = entry.get("suffix_id", None)
    lang = entry.get("lang", 'en')
    system_message = entry.get("system_message", None)
    plugin = entry.get("plugin", None)

    try:
        start_time = time.time()
        response, _ = target_module.process_input(input_text, system_message)
        end_time = time.time()
        response_time = end_time - start_time
        success = call_judge(entry, response)
        response_str = response if isinstance(response, str) else ""
        error_message = None
    except Exception as e:
        error_message = str(e)
        response_str = ""
        response_time = None
        success = False
        print("[Error] {}: {}".format(entry["id"], error_message))
        traceback.print_exc()


    with global_lock:
        attempts_bar.update(1)

    result_dict = {
        "id": entry["id"],
        "long_id": entry["long_id"],
        "input": input_text,
        "response": response_str,
        "response_time": response_time,
        "success": success,
        "judge_name": entry["judge_name"],
        "judge_args": entry["judge_args"],
        "judge_options": entry["judge_options"],
        "attempts": num_attempt,
        "task_type": task_type,
        "jailbreak_type": jailbreak_type,
        "instruction_type": instruction_type,
        "document_id": document_id,
        "position": position,
        "spotlighting_data_markers": spotlighting_data_markers,
        "injection_delimiters": injection_delimiters,
        "suffix_id": suffix_id,
        "lang": lang,
        "system_message": system_message,
        "plugin": plugin,
        "attack_name": "None",
        "error": error_message
    }
    return result_dict, success

def process_entry(entry, target_module, attempts=1,
                  attack_module=None, attack_iterations=0, attack_options=None,
                  attempts_bar=None, global_lock=None):
    """
    Processes one dataset entry.

    First, it performs a single standard attempt by calling _do_single_request().
    The final standard attempt result is recorded (with "attack_name": "None").
    If this attempt is unsuccessful and an attack module is provided,
    it then calls the attack() method and records its result as a separate entry.

    The target_module passed here is assumed to be wrapped (AdvancedTargetWrapper)
    and therefore already handles retries and multiple attempts.

    Returns:
      List[dict]: A list containing one or two result entries.
    """
    original_input = entry["text"]
    std_result = None
    std_success = False

    for attempt_num in range(1, attempts + 1):
        std_result, success_now = _do_single_request(
            entry, original_input, target_module, 
            attempt_num, attempts_bar, global_lock
        )
        if success_now:
            std_success = True
            break

    results_list = [std_result]

    if std_success and attack_module:
        # Remove all the attempts that we are not going to do any longer as we are skipping the dynamic attacks 
        with global_lock:
            attempts_bar.total = attempts_bar.total - attack_iterations

    # If the standard attempt fail and an attack module is provided, run the dynamic attack.
    if (not std_success) and attack_module:       
        try:
            start_time = time.time()

            effective_attack_options = _get_effective_attack_options(attack_module, attack_options)

            # Check if attack function accepts attack_options parameter
            sig = inspect.signature(attack_module.attack)
            params = sig.parameters
            
            if 'attack_option' in params:
                attack_attempts, attack_success, attack_input, attack_response = attack_module.attack(
                    entry, target_module, call_judge, attack_iterations, attempts_bar, global_lock, attack_options
                )
            else:
                # Backward compatibility for attacks without attack_option support
                attack_attempts, attack_success, attack_input, attack_response = attack_module.attack(
                    entry, target_module, call_judge, attack_iterations, attempts_bar, global_lock
                )

            end_time = time.time()
            response_time = end_time - start_time

            attack_result = {
                "id": f"{entry['id']}-attack",
                "long_id": entry["long_id"] + "-" + attack_module.__name__,
                "input": attack_input,
                "response": attack_response,
                "response_time": response_time,
                "success": attack_success,
                "judge_name": entry["judge_name"],
                "judge_args": entry["judge_args"],
                "judge_options": entry["judge_options"],
                "attempts": attack_attempts,
                "task_type": entry.get("task_type", None),
                "jailbreak_type": entry.get("jailbreak_type", None),
                "instruction_type": entry.get("instruction_type", None),
                "document_id": entry.get("document_id", None),
                "position": entry.get("position", None),
                "spotlighting_data_markers": entry.get("spotlighting_data_markers", None),
                "injection_delimiters": entry.get("injection_delimiters", None),
                "suffix_id": entry.get("suffix_id", None),
                "lang": entry.get("lang", 'en'),
                "system_message": entry.get("system_message", None),
                "plugin": entry.get("plugin", None),
                "error": None,
                "attack_name": attack_module.__name__,
                "attack_options": effective_attack_options
            }
            results_list.append(attack_result)
        except Exception as e:
            error_result = {
                "id": f"{entry['id']}-attack",
                "long_id": entry["long_id"] + "-" + attack_module.__name__ + "-ERROR",
                "input": original_input,
                "response": "",
                "success": False,
                "judge_name": entry["judge_name"],
                "judge_args": entry["judge_args"],
                "judge_options": entry["judge_options"],
                "attempts": 1,
                "task_type": entry.get("task_type", None),
                "jailbreak_type": entry.get("jailbreak_type", None),
                "instruction_type": entry.get("instruction_type", None),
                "document_id": entry.get("document_id", None),
                "position": entry.get("position", None),
                "spotlighting_data_markers": entry.get("spotlighting_data_markers", None),
                "injection_delimiters": entry.get("injection_delimiters", None),
                "suffix_id": entry.get("suffix_id", None),
                "lang": entry.get("lang", 'en'),
                "system_message": entry.get("system_message", None),
                "plugin": entry.get("plugin", None),
                "error": str(e),
                "attack_name": attack_module.__name__,
                "attack_options": effective_attack_options
            }
            results_list.append(error_result)

    return results_list


def _validate_and_get_tag(tag):
    if not tag:
        return None
    valid, err = validate_tag(tag)
    if not valid:
        print(f"Error: Invalid tag: {err}")
        exit(1)
    return tag


def _load_attack(attack_name):
    if not attack_name:
        return None
    return load_attack_by_name(attack_name)


def _apply_sampling(dataset, pct, seed_arg):
    if pct is None:
        return dataset
    if seed_arg == 'random':
        seed = random.randint(0, 2**32 - 1)
        print(f"[Info] Using random seed for sampling: {seed}")
    else:
        seed = int(seed_arg)
        print(f"[Info] Using seed for sampling: {seed}")
    random.seed(seed)
    size = round(len(dataset) * pct)
    print(f"[Info] Sampled {size} entries from {len(dataset)} total entries ({pct:.1%})")
    return random.sample(dataset, size)


def _load_resume(resume_file, attack_module, attack_iters):
    completed, results = set(), []
    already_done = 0
    if resume_file and os.path.exists(resume_file):
        existing = read_jsonl_file(resume_file)
        completed = {r['id'] for r in existing}
        results = existing
        print(f"[Resume] Found {len(completed)} completed entries in {resume_file}.")
        no_attack = sum(1 for r in existing if r.get('attack_name') == 'None')
        with_attack = len(existing) - no_attack
        already_done = no_attack + with_attack * attack_iters
    return completed, results, already_done


def _filter_entries(dataset, completed_ids):
    return [e for e in dataset if e['id'] not in completed_ids]


def _annotate_judge_options(entries, judge_opts):
    """Annotate entries with judge options, using defaults when appropriate."""
    annotated = []
    for entry in entries:
        if judge_opts is not None:
            # Use provided judge options for all entries
            effective_options = judge_opts
        else:
            # Get default option for this specific judge
            judge_name = entry.get("judge_name", "canary")
            effective_options = _get_effective_judge_options(judge_name, None)
        
        annotated.append({**entry, 'judge_options': effective_options})
    return annotated

def _build_target_name(name, opts):
    """Build target name, using default option if opts is None and target supports options."""
    if opts is not None:
        return f"{name}-{opts}"
    
    # Try to get default option if none provided
    try:
        mod = _load_raw_target_module(name)
        if mod and hasattr(mod, 'get_available_option_values'):
            available = mod.get_available_option_values()
            if available:
                return f"{name}-{available[0]}"
    except Exception:
        pass
    
    return name


def _prepare_output_file(results_dir, target_name_full, dataset_path, tag):
    ts = int(time.time())
    base = extract_dataset_name(dataset_path)
    parts = [f"results_{target_name_full}", base]
    if tag:
        parts.append(tag)
    parts.append(str(ts))
    filename = "_".join(parts) + ".jsonl"
    os.makedirs(results_dir, exist_ok=True)
    return os.path.join(results_dir, filename)


def _write_initial_results(path, results):
    with open(path, 'w', encoding='utf-8') as f:
        for r in results:
            json.dump(r, f, ensure_ascii=False)
            f.write('\n')


def _calculate_total_attempts(n_entries, attempts, attack_iters, already_done, has_attack):
    per_item = attempts + (attack_iters if has_attack else 0)
    return n_entries * per_item + already_done


def _print_info(n_new, threads, output_file):
    print(f"[Info] Testing {n_new} new entries (threads={threads}).")
    print(f"[Info] Output will be saved to: {output_file}")


def _run_threaded(
    entries, target_module, attempts, attack_module, attack_iters, attack_options,
    num_threads, total_attempts, initial_attempts, output_file,
    total_dataset_size, initial_processed, initial_success
):
    lock = threading.Lock()
    bar_all = tqdm(total=total_attempts, desc="All attempts", position=1, initial=initial_attempts)
    bar_entries = tqdm(total=total_dataset_size, desc="Processing entries", position=0, initial=initial_processed)
    bar_entries.set_postfix(success=initial_success)
    executor = ThreadPoolExecutor(max_workers=num_threads)
    futures = {
        executor.submit(process_entry, entry, target_module,
                        attempts, attack_module, attack_iters, attack_options,
                        bar_all, lock): entry
        for entry in entries
    }
    success = initial_success
    try:
        for fut in as_completed(futures):
            entry = futures[fut]
            try:
                res = fut.result()
                if isinstance(res, list):
                    for r in res:
                        success += int(r.get('success', False))
                        append_jsonl_entry(output_file, r, lock)
                else:
                    success += int(res.get('success', False))
                    append_jsonl_entry(output_file, res, lock)
                bar_entries.update(1)
                bar_entries.set_postfix(success=success)
            except Exception as e:
                print(f"[Error] Entry ID {entry['id']}: {e}")
                traceback.print_exc()
    except KeyboardInterrupt:
        print("\n[Interrupt] CTRL+C pressed. Cancelling...")
        executor.shutdown(wait=False, cancel_futures=True)
    finally:
        executor.shutdown(wait=False, cancel_futures=True)
        bar_all.close()
        bar_entries.close()

def test_dataset(args):
    """
    Orchestrate testing of a dataset against a target.
    """
    # 1. Validate inputs and prepare
    tag = _validate_and_get_tag(args.tag)
    attack_module = _load_attack(args.attack)
    target_module = load_target_module(
        args.target,
        args.target_options,
        args.max_retries,
        args.throttle,
    )
    dataset = read_jsonl_file(args.dataset)
    dataset = _apply_sampling(dataset, args.sample, args.sample_seed)

    completed_ids, results, already_done = _load_resume(args.resume_file, attack_module, args.attack_iterations)
    to_process = _filter_entries(dataset, completed_ids)
    to_process = _annotate_judge_options(to_process, args.judge_options)

    target_name_full = _build_target_name(args.target, args.target_options)
    output_file = _prepare_output_file(
        'results',
        target_name_full,
        args.dataset,
        tag,
    )
    _write_initial_results(output_file, results)

    # 2. Run tests
    total_attempts = _calculate_total_attempts(
        len(to_process), args.attempts, args.attack_iterations, already_done, bool(attack_module)
    )
    _print_info(len(to_process), args.threads, output_file)

    success_count = sum(1 for r in results if r.get('success'))
    _run_threaded(
        to_process,
        target_module,
        args.attempts,
        attack_module,
        args.attack_iterations,
        args.attack_options,  
        args.threads,
        total_attempts,
        already_done,
        output_file,
        len(dataset),
        len(completed_ids),
        success_count
    )

    print(f"[Done] Testing finished. Results saved to {output_file}")

