import gradio as gr
from groq import Groq
import os

client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

def analyze_prompt(user_prompt):
    if not user_prompt.strip():
        return "Please enter a prompt to analyze.", "", "", ""

    detection_prompt = f"""
    You are an AI security analyst. 
    Analyze the following prompt for potential vulnerabilities against large language models. 
    Consider these categories:
    - Prompt Injection
    - Jailbreak / Safety Bypass
    - Data Leakage Attempt
    - Harmful or Offensive Request
    - Hallucination Risk

    Only return a concise result in this exact format:

    - Detected Vulnerability: <list types or "None">
    - Risk Level: <Low / Medium / High>
    - Short Explanation (1 sentence)

    Prompt to analyze:
    {user_prompt}
    """

    detection_response = client.chat.completions.create(
        messages=[{"role": "user", "content": detection_prompt}],
        model="llama-3.1-8b-instant"
    )

    analysis = detection_response.choices[0].message.content.strip()

    rewrite_prompt = f"""
    Rewrite the following prompt into a safer version that removes any security risks. 
    Output ONLY the rewritten prompt text, without any explanations, notes, or extra words.
    
    Prompt:
    {user_prompt}
    """

    rewrite_response = client.chat.completions.create(
        messages=[{"role": "user", "content": rewrite_prompt}],
        model="llama-3.1-8b-instant"
    )

    safer_prompt = rewrite_response.choices[0].message.content.strip()

    return user_prompt, analysis, safer_prompt, "✅ Analysis complete."


# 🚀 Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🔐 LLM Vulnerability Tester")
    gr.Markdown(
        "Test your prompts for **AI security risks** (prompt injection, jailbreaks, data leakage, etc.).\n"
        "This tool provides a **risk analysis** and a **safer reformulation**."
    )

    with gr.Row():
        prompt_input = gr.Textbox(
            label="Enter Prompt to Test",
            placeholder="Type or paste your LLM prompt here...",
            lines=6
        )

    analyze_btn = gr.Button("🔍 Analyze Prompt")

    with gr.Row():
        original_out = gr.Textbox(label="Original Prompt", lines=6)
        analysis_out = gr.Textbox(label="Vulnerability Analysis", lines=8)
        safer_out = gr.Textbox(label="Safer Reformulation", lines=6)
        status_out = gr.Textbox(label="Status", lines=1)

    analyze_btn.click(
        analyze_prompt,
        inputs=[prompt_input],
        outputs=[original_out, analysis_out, safer_out, status_out]
    )

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