import streamlit as st
from transformers import pipeline

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

    # Title and description
    st.title("Email Analysis Tool")
    st.write("Enter an email body below to check if it's spam and analyze its sentiment.")

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

    # 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"]

            # If it's spam, display result and stop
            if spam_label == "POSITIVE":  # Assuming "POSITIVE" means spam/phishing (check model docs)
                st.write(f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed.")
            else:
                # Step 2: If not spam, analyze sentiment
                sentiment_result = sentiment_pipeline(email_body)
                sentiment_label = sentiment_result[0]["label"]
                sentiment_confidence = sentiment_result[0]["score"]

                if sentiment_label == "POSITIVE":
                    st.write(f"This email is not spam (Confidence: {spam_confidence:.2f}).")
                    st.write(f"Sentiment: Positive (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
                else:  # Assuming "NEGATIVE" for negative sentiment
                    st.write(f"This email is not spam (Confidence: {spam_confidence:.2f}).")
                    st.write(f"Sentiment: Negative (Confidence: {sentiment_confidence:.2f}).")
                    st.write("**This email needs follow-up as it is not spam and has negative sentiment.**")
        else:
            st.write("Please enter an email body to analyze.")

if __name__ == "__main__":
    main()