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

# Function to analyze email for spam and sentiment
def analyze_email(email_body):
    # Load pre-trained models: spam detection and sentiment analysis
    spam_pipeline = pipeline("text-classification", model="cybersectony/phishing-email-detection-distilbert_v2.4.1")
    sentiment_model = AutoModelForSequenceClassification.from_pretrained("ISOM5240GP4/email_sentiment", num_labels=2)
    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

    # Step 1: Check if the email is spam using the spam detection model
    spam_result = spam_pipeline(email_body)
    spam_label = spam_result[0]["label"]  # LABEL_1 indicates spam
    spam_confidence = spam_result[0]["score"]

    if spam_label == "LABEL_1":
        # If spam, return type "spam" and a message indicating no follow-up
        return "spam", f"This is a spam email (Confidence: {spam_confidence:.2f}). No follow-up needed."
    else:
        # Step 2: For non-spam emails, analyze sentiment (positive/negative)
        inputs = tokenizer(email_body, padding=True, truncation=True, return_tensors='pt')  # Tokenize input
        outputs = sentiment_model(**inputs)  # Get model predictions
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)  # Apply softmax for probabilities
        predictions = predictions.cpu().detach().numpy()  # Convert to numpy array
        sentiment_index = np.argmax(predictions)  # Get the predicted sentiment (0 = negative, 1 = positive)
        sentiment_confidence = predictions[0][sentiment_index]
        sentiment = "Positive" if sentiment_index == 1 else "Negative"

        if sentiment == "Positive":
            # If positive sentiment, no follow-up needed
            return "positive", (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
                                f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}). No follow-up needed.")
        else:
            # If negative sentiment, mark as needing follow-up with bolded text
            return "negative", (f"This email is not spam (Confidence: {spam_confidence:.2f}).\n"
                                f"Sentiment: {sentiment} (Confidence: {sentiment_confidence:.2f}).\n"
                                "<b>Need to Follow-Up</b>: This email is not spam and has negative sentiment.")

# Main application function
def main():
    # Set the app title to the project name
    st.title("EmailSentry")
    # Display the project objective
    st.write("Aims to perform analysis on incoming emails and to determine whether there is urgency or higher priority for the company to follow-up.")

    # Initialize session state to store email input and analysis results
    if "email_body" not in st.session_state:
        st.session_state.email_body = ""  # Holds the email text
    if "result" not in st.session_state:
        st.session_state.result = ""  # Stores the analysis result text
    if "result_type" not in st.session_state:
        st.session_state.result_type = ""  # Tracks result type (spam, positive, negative)

    # Add collapsible instructions for user guidance
    with st.expander("How to Use", expanded=False):
        st.write("""
        - Type or paste an email into the text box.
        - Alternatively, click one of the sample buttons to load a predefined email.
        - Press 'Analyze Email' to check if it’s spam and analyze its sentiment.
        - Use 'Clear' to reset the input and result.
        """)

    # Text area where users input or view the email body
    email_body = st.text_area("Email Body", value=st.session_state.email_body, height=200, key="email_input")

    # Define sample emails and their snippets for example buttons
    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
    """
    spam_snippet = "Subject: Urgent: Verify Your Account Now! Dear Customer, We have detected unusual activity..."

    sample_not_spam_positive = """
Subject: Great Experience with HKTV mall
Dear Sir,
I just received my order and I’m really impressed with the speed of the delivery. Keep up the good work.
Best regards,
Emily
    """
    positive_snippet = "Subject: Great Experience with HKTV mall Dear Sir, I just received my order and I’m really..."

    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
    """
    negative_snippet = "Subject: Issue with Recent Delivery Dear Support, I received my package today, but..."

    # Custom CSS for styling buttons and result boxes with higher specificity
    st.markdown("""
        <style>
        /* Style for sample buttons (smaller, light gray), excluding Clear */
        div.stButton > button[kind="secondary"]:not([key="clear"]) {
            font-size: 12px !important;
            padding: 5px 10px !important;
            background-color: #f0f0f0 !important;
            color: #333333 !important;
            border: 1px solid #cccccc !important;
            border-radius: 3px !important;
        }
        /* Analyze Email button (force orange, matches Clear button size) */
        div.stButton > button[key="analyze"] {
            background-color: #FF5733 !important; /* Force orange color */
            color: white !important; /* Force white text */
            font-size: 18px !important;
            padding: 12px 24px !important;
            border: none !important;
            border-radius: 5px !important;
            width: 100% !important;
            height: 50px !important;
            box-sizing: border-box !important;
            text-align: center !important;
            margin: 0 !important; /* Remove any default margins */
        }
        div.stButton > button[key="analyze"]:hover {
            background-color: #E74C3C !important; /* Darker orange on hover */
        }
        /* Clear button (gray, aligned with Analyze) */
        div.stButton > button[key="clear"] {
            background-color: #d3d3d3 !important; /* Force gray color */
            color: #333333 !important; /* Force dark text */
            font-size: 18px !important;
            padding: 12px 24px !important;
            border: none !important;
            border-radius: 5px !important;
            width: 100% !important;
            height: 50px !important;
            box-sizing: border-box !important;
            text-align: center !important;
            margin: 0 !important; /* Remove any default margins */
        }
        div.stButton > button[key="clear"]:hover {
            background-color: #b0b0b0 !important; /* Darker gray on hover */
        }
        /* Result boxes with updated colors */
        .spam-result {
            background-color: #ff3333 !important; /* Red for no follow-up (spam) */
            color: white !important;
            padding: 10px !important;
            border-radius: 5px !important;
            border: 1px solid #cc0000 !important;
        }
        .positive-result {
            background-color: #ff3333 !important; /* Red for no follow-up (positive) */
            color: white !important;
            padding: 10px !important;
            border-radius: 5px !important;
            border: 1px solid #cc0000 !important;
        }
        .negative-result {
            background-color: #006633 !important; /* Dark green for follow-up needed */
            color: white !important;
            padding: 10px !important;
            border-radius: 5px !important;
            border: 1px solid #004d26 !important;
        }
        </style>
    """, unsafe_allow_html=True)

    # Subheading to label the sample email buttons
    st.subheader("Examples")

    # Layout for sample buttons in 3 columns
    col1, col2, col3 = st.columns(3)
    with col1:
        # Button to load spam sample
        if st.button(spam_snippet, key="spam_sample"):
            st.session_state.email_body = sample_spam
            st.session_state.result = ""
            st.session_state.result_type = ""
            st.rerun()
    with col2:
        # Button to load positive non-spam sample
        if st.button(positive_snippet, key="positive_sample"):
            st.session_state.email_body = sample_not_spam_positive
            st.session_state.result = ""
            st.session_state.result_type = ""
            st.rerun()
    with col3:
        # Button to load negative non-spam sample
        if st.button(negative_snippet, key="negative_sample"):
            st.session_state.email_body = sample_not_spam_negative
            st.session_state.result = ""
            st.session_state.result_type = ""
            st.rerun()

    # Layout for action buttons (Analyze and Clear) in 2 columns
    col_analyze, col_clear = st.columns(2)
    with col_analyze:
        # Button to trigger email analysis (no type="primary" to rely on CSS)
        if st.button("Analyze Email", key="analyze"):
            if email_body:
                with st.spinner("Analyzing email..."):  # Show spinner during processing
                    result_type, result = analyze_email(email_body)
                    st.session_state.result = result
                    st.session_state.result_type = result_type
            else:
                # Error message if no email is provided
                st.session_state.result = "Please enter an email body or select a sample to analyze."
                st.session_state.result_type = ""
    with col_clear:
        # Button to reset the app state
        if st.button("Clear", key="clear"):
            st.session_state.email_body = ""
            st.session_state.result = ""
            st.session_state.result_type = ""
            st.rerun()

    # Display the analysis result in styled boxes based on result type
    if st.session_state.result:
        if st.session_state.result_type == "spam":
            st.markdown(f'<div class="spam-result">{st.session_state.result}</div>', unsafe_allow_html=True)
        elif st.session_state.result_type == "positive":
            st.markdown(f'<div class="positive-result">{st.session_state.result}</div>', unsafe_allow_html=True)
        elif st.session_state.result_type == "negative":
            st.markdown(f'<div class="negative-result">{st.session_state.result}</div>', unsafe_allow_html=True)
        else:
            st.write(st.session_state.result)  # Display error messages without styling

# Run the app
if __name__ == "__main__":
    main()