import streamlit as st
import torch
from transformers import BertTokenizer, BertForSequenceClassification

# ✅ Ensure set_page_config is the first Streamlit command
st.set_page_config(page_title="Scam Slayer", layout="centered")

# Load model from Hugging Face
MODEL_NAME = "sellestas/scam_slayer_model"

try:
    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
    model = BertForSequenceClassification.from_pretrained(MODEL_NAME)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    st.success("✅ Scam Slayer Model Loaded Successfully!")
except Exception as e:
    st.error(f"❌ Error loading model: {e}")

# Function to classify email
def classify_email(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
        confidence, prediction = torch.max(probabilities, dim=-1)

    label_map = {0: "Non-Malicious ✅", 1: "Malicious 🚨"}
    return label_map[prediction.item()], confidence.item() * 100

# UI Layout
st.image("logo.png", width=150)
st.title("🛡️ Scam Slayer - AI Email Threat Detector")
st.markdown("### 🔍 Detect phishing and malicious emails instantly!")

# Sidebar About Button
with st.sidebar:
    if st.button("ℹ️ About Scam Slayer"):
        st.markdown("""
        ## 📌 About Scam Slayer  
        **AI-powered cybersecurity tool** to detect phishing threats.  

        ✅ **Purpose**: Identify and stop phishing attacks.  
        ✅ **Model**: Fine-tuned BERT-based classifier.  
        ✅ **Developed for**: **SANS AI Cybersecurity Hackathon 2025**.  
        ✅ **Features**:  
        - Detects **Malicious & Non-Malicious** emails  
        - Uses **NLP** for content analysis  
        - Provides a **confidence score** (1-100%)  

        **Version**: 1.0.0  
        """)

# Email Input
email_text = st.text_area("✉️ Paste the email content below:", height=200)

# Detect Scam Button
if st.button("🚀 Detect Scam", help="Click to analyze the email content"):
    if email_text.strip():
        category, confidence = classify_email(email_text)
        st.success(f"**🔹 Result: {category} ({confidence:.2f}% Confidence)**")
        st.markdown("✅ **Stay vigilant against scams!** 🚀")
    else:
        st.warning("⚠️ Please enter email content to analyze!")