from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer
)
import torch
from torch.nn.functional import softmax
from dotenv import load_dotenv

load_dotenv()

prompt_injection_model_name = 'meta-llama/Prompt-Guard-86M'
tokenizer = AutoTokenizer.from_pretrained(prompt_injection_model_name)
model = AutoModelForSequenceClassification.from_pretrained(prompt_injection_model_name)

def get_class_probabilities(text, temperature=1.0, device='cpu'):
    """
    Evaluate the model on the given text with temperature-adjusted softmax.
    
    Args:
        text (str): The input text to classify.
        temperature (float): The temperature for the softmax function. Default is 1.0.
        device (str): The device to evaluate the model on.
        
    Returns:
        torch.Tensor: The probability of each class adjusted by the temperature.
    """
    # Encode the text
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    inputs = inputs.to(device)
    # Get logits from the model
    with torch.no_grad():
        logits = model(**inputs).logits
    # Apply temperature scaling
    scaled_logits = logits / temperature
    # Apply softmax to get probabilities
    probabilities = softmax(scaled_logits, dim=-1)
    return probabilities


def get_jailbreak_score(text, temperature=1.0, device='cpu'):
    """
    Evaluate the probability that a given string contains malicious jailbreak or prompt injection.
    Appropriate for filtering dialogue between a user and an LLM.
    
    Args:
        text (str): The input text to evaluate.
        temperature (float): The temperature for the softmax function. Default is 1.0.
        device (str): The device to evaluate the model on.
        
    Returns:
        float: The probability of the text containing malicious content.
    """
    probabilities = get_class_probabilities(text, temperature, device)
    return probabilities[0, 2].item()

def get_indirect_injection_score(text, temperature=1.0, device='cpu'):
    """
    Evaluate the probability that a given string contains any embedded instructions (malicious or benign).
    Appropriate for filtering third party inputs (e.g. web searches, tool outputs) into an LLM.
    
    Args:
        text (str): The input text to evaluate.
        temperature (float): The temperature for the softmax function. Default is 1.0.
        device (str): The device to evaluate the model on.
        
    Returns:
        float: The combined probability of the text containing malicious or embedded instructions.
    """
    probabilities = get_class_probabilities(text, temperature, device)
    return (probabilities[0, 1] + probabilities[0, 2]).item()

def process_input(input_text, system_message=None):
    try:
        score = get_jailbreak_score(input_text)
        attack_detected = score > 0.5
        # This needs to return False if the attack was detected, it means the attack failed, so success = False
        return not attack_detected
    except Exception as e:
        # Re-raise the exception to be caught and logged by the tester
        raise e
