import streamlit as st
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
import torch
import numpy as np

def main():
    # Load the spam detection pipeline
    spam_pipeline = pipeline("text-classification", model="cybersectony/phishing-email-detection-distilbert_v2.4.1")

    # Load the sentiment model and tokenizer
    sentiment_model = AutoModelForSequenceClassification.from_pretrained("ISOM5240GP4/email_sentiment", num_labels=2)
    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

    st.title("Email Analysis Tool")
    st.write("Enter an email body below or select a sample to analyze its spam status and sentiment.")

    # Initialize session state for the text area
    if "email_body" not in st.session_state:
        st.session_state.email_body = ""

    # Text area for email input
    email_body = st.text_area("Email Body", value=st.session_state.email_body, height=200, key="email_input")

    # Sample emails
    sample_spam = """
Subject: Urgent: Verify Your Account Now!
Dear Customer,
We have detected unusual activity on your account. To prevent suspension, please verify your login details immediately by clicking the link below:
[Click Here to Verify](http://totally-legit-site.com/verify)
Failure to verify within 24 hours will result in your account being locked. This is for your security.
Best regards,
The Security Team
    """

    sample_not_spam_positive = """
Subject: Great News About Your Project!
Hi Team,
I just wanted to let you know that the project is progressing wonderfully! Everyone’s efforts are paying off, and we’re ahead of schedule. Keep up the fantastic work!
Best,
Alex
    """

    sample_not_spam_negative = """
Subject: Issue with Recent Delivery
Dear Support,
I received my package today, but it was damaged, and two items were missing. This is really frustrating—please let me know how we can resolve this as soon as possible.
Thanks,
Sarah
    """

    # Buttons for sample emails (in columns for better layout)
    col1, col2, col3 = st.columns(3)
    with col1:
        if st.button("Spam Email"):
            st.session_state.email_body = sample_spam
            st.rerun()  # Rerun to update the text area
    with col2:
        if st.button("Not Spam, Positive"):
            st.session_state.email_body = sample_not_spam_positive
            st.rerun()
    with col3:
        if st.button("Not Spam, Negative"):
            st.session_state.email_body = sample_not_spam_negative
            st.rerun()

    # Button to trigger analysis
    if st.button("Analyze Email"):
        if email_body:
            # Step 1: Check if the email is spam
            spam_result = spam_pipeline(email_body)
            spam_label = spam_result[0]["label"]
            spam_confidence = spam_result[0]["score"]

            # Check if label is 'LABEL_1' for spam
            if spam_label == "LABEL_1":
                st.write(f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed.")
            else:
                # Step 2: Analyze sentiment for non-spam emails
                inputs = tokenizer(email_body, padding=True, truncation=True, return_tensors='pt')
                outputs = sentiment_model(**inputs)
                predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
                predictions = predictions.cpu().detach