import gradio as gr
from huggingface_hub import InferenceClient
from transformers import pipeline
from typing import List, Dict, Any

def generate_attack(
    prompt: str,
    history: List[Any],
) -> List[str]:
    """
    Simulates a Blackhat AI scenario by generating attack strategies and potential impacts.
    
    Args:
        prompt (str): The user's input to the simulator.
        history (List[Dict]): The user's message history with timestamps.
    
    Returns:
        List[str]: A list of attack responses from the AI.
    """
    messages = [
        {
            "role": "system",
            "content": f"Responding to {prompt}..."
        }
    ]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": prompt})
    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=100,  # limit the length of responses
        stream=True,
        temperature=0.5,  # adjust the temperature to improve response quality
        top_p=None,      # disable top-p filtering
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

def simulate_attack(
    prompt: str,
    history: List[Any],
) -> List[str]:
    """
    Simulates a Blackhat AI scenario by generating attack strategies and potential impacts.
    
    Args:
        prompt (str): The user's input to the simulator.
        history (List[Dict]): The user's message history with timestamps.
    
    Returns:
        List[str]: A list of attack responses from the AI.
    """
    messages = [
        {
            "role": "system",
            "content": f"Simulating a Blackhat AI scenario..."
        }
    ]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": prompt})
    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=100,  # limit the length of responses
        stream=True,
        temperature=0.5,  # adjust the temperature to improve response quality
        top_p=None,      # disable top-p filtering
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

# Define the Gradio ChatInterface with security-focused configuration
demo = gr.ChatInterface(
    generate_attack,
    additional_inputs=[
        gr.Textbox(
            value=(
                "You are an AI simulator for cybersecurity training, designed to generate attack scenarios, analyze their impacts, and suggest countermeasures."
            ),
            label="System message",
        ),
        # ... (other inputs)
    ],
    title="Blackhat AI Simulator",
    description=(
        "This simulator generates adversarial scenarios, analyzes attack vectors, and provides ethical countermeasures. Use responsibly for cybersecurity training and awareness."
    ),
)

if __name__ == "__main__":
    demo.launch()