# streamlit_app.py
import streamlit as st
from models import load_model
from utils import mask_pii

st.set_page_config(page_title="📧 Email Classification System", page_icon="📧")

st.title("📧 Email Classification and PII Masking System")

# Load model
model = load_model("model/classifier.pkl")

# Prediction function
def classify_email_direct(email_body):
    masked_text, entities = mask_pii(email_body)
    category = model.predict([masked_text])[0]
    return {
        "input_email_body": email_body,
        "list_of_masked_entities": entities,
        "masked_email": masked_text,
        "category_of_the_email": category
    }

# Streamlit UI
email_input = st.text_area("✉️ Enter your email text:")

if st.button("🚀 Classify Email"):
    if email_input.strip() != "":
        result = classify_email_direct(email_input)

        st.subheader("🔎 Masked Email")
        st.write(result['masked_email'])

        st.subheader("🔐 List of Masked Entities")
        st.json(result['list_of_masked_entities'])

        st.subheader("📂 Predicted Category")
        st.success(result['category_of_the_email'])
    else:
        st.warning("⚠️ Please enter some text first.")