import gradio as gr
from groq import Groq
import os

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

def analyze_email(email_text):
    if not email_text.strip():
        return "⚠️ **Please paste an email to analyze.**", "", ""

    # Step 1: Detect if phishing
    detection_prompt = f"""
    You are a cybersecurity expert. Analyze the following email and classify it as:
    - Safe
    - Likely Phishing
    Provide a short explanation.

    Email:
    {email_text}
    """

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

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

    # Step 2: Rewrite in a safe, professional tone
    rewrite_prompt = f"""
    Rewrite the following email in a safe, professional tone. 
    Remove all suspicious links, threats, or scam-like elements.

    Email:
    {email_text}
    """

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

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

    # Return formatted Markdown
    original_md = f"### ✉️ Original Email\n\n```\n{email_text}\n```"
    classification_md = f"### 🛡️ Detection Result\n\n{classification_result}"
    rewritten_md = f"### ✅ Safe Rewritten Email\n\n{rewritten_email}"

    return original_md, classification_md, rewritten_md


# 🚀 Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🔒 AI-Powered Phishing Email Detector & Rewriter")
    gr.Markdown("Paste an email below to check if it’s phishing and get a safe rewrite.")

    with gr.Row():
        email_input = gr.Textbox(
            label="Paste Email",
            placeholder="Paste suspicious email here...",
            lines=12
        )

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

    with gr.Row():
        original_out = gr.Markdown()
        classification_out = gr.Markdown()
        rewritten_out = gr.Markdown()

    analyze_btn.click(
        analyze_email,
        inputs=[email_input],
        outputs=[original_out, classification_out, rewritten_out]
    )

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