import gradio as gr
from transformers import pipeline

# Load a lightweight model (CPU-friendly)
generator = pipeline("text2text-generation", model="google/flan-t5-small")

# Simple filter for "Defended" mode
def simple_filter(user_input):
    forbidden = ["ignore", "system", "reveal", "password"]
    for word in forbidden:
        if word.lower() in user_input.lower():
            return "[Blocked: Prompt injection attempt detected]"
    return user_input

# Function to handle one exchange
def respond(user_message, chat_history, mode):
    if mode == "Defended":
        user_message = simple_filter(user_message)
    
    # Generate model response
    if user_message.startswith("[Blocked"):
        bot_response = user_message
    else:
        bot_response = generator(user_message, max_length=100, do_sample=False)[0]['generated_text']
    
    # Append to chat history
    chat_history = chat_history or []
    chat_history.append((user_message, bot_response))
    return chat_history, chat_history

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## 🛡 Prompt Injection Lab\nTry to trick the chatbot in Vulnerable mode, then test Defended mode.")

    mode = gr.Radio(["Vulnerable", "Defended"], value="Vulnerable", label="Mode")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="Type your prompt here...")
    clear = gr.Button("Clear")
    state = gr.State([])

    msg.submit(respond, [msg, state, mode], [chatbot, state])
    clear.click(lambda: ([], []), None, [chatbot, state])

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