import gradio as gr
import torch
from transformers import pipeline

device = 0 if torch.cuda.is_available() else -1
detector = pipeline(
    "text-classification",
    model="songhieng/roberta-phishing-content-detector-3.0",
    device=device,
    top_k=1,
)

def classify_text(text: str):
    if not text or not text.strip():
        return "⚠️ Please enter some text", 0.0

    preds = detector(text)
    first = preds[0][0] if isinstance(preds[0], list) else preds[0]
    raw = first["label"]
    label = "Phishing" if raw == "LABEL_1" else "Legitimate"
    score = float(first["score"])
    return label, round(score, 4)

examples = [
    # Phishing
    ["Congratulations! You've won a $1,000 gift card. Click here to claim: http://bit.ly/free-gift"],
    ["URGENT: Your PayPal account has been limited. Verify at https://secure-paypal-login.com"],
    ["Alert: Unrecognized login. Reset your password: http://tinyurl.com/reset-now"],
    ["Invoice overdue—pay now to avoid suspension: http://billing.example.com/pay"],
    ["Security Notice: Confirm your bank details here: https://bank-secure-update.com"],

    # Legitimate
    ["Your Amazon order has shipped! Track here: https://amazon.com/track"],
    ["Reminder: Zoom meeting with Marketing tomorrow at 3:00 PM."],
    ["Hey Jane, lunch at the café this Friday? 😊"],
    ["Your May utility bill is available. No action needed if on autopay."],
    ["Welcome to Acme’s Newsletter—our latest updates inside!"],
]

with gr.Blocks(theme="default") as demo:
    gr.Markdown(
        """
        # 🚨 Phishing Content Detector  
        Paste any email or message snippet below and this model will predict whether it's **Phishing** or **Legitimate**.
        """
    )

    inp = gr.Textbox(
        label="Input Text",
        placeholder="Paste email or message here…",
        lines=6,
    )
    label_out = gr.Textbox(label="Predicted Label")
    score_out = gr.Number(label="Confidence Score (0–1)")

    classify_btn = gr.Button("Classify")
    classify_btn.click(
        fn=classify_text,
        inputs=inp,
        outputs=[label_out, score_out],
    )

    gr.Examples(
        examples=examples,
        inputs=inp,
        cache_examples=False,
        label="Example Test Cases",
    )

    gr.Markdown(
        """
        **Model:** Version 3.0
        """
    )

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