# app.py
import streamlit as st
from transformers import pipeline
from llama import load_llama_model, generate_llama_summary, PROMPT_TEMPLATE

st.set_page_config(page_title="Email Summarizer", layout="wide")

# Predefined examples with reference summaries
EXAMPLES = {
    "Sample1": {
        "email": """
        
Sehr geehrte Damen und Herren, 
Leider wurden wir von einem Unwetter getroffen, weshalb unser Keller vollgelaufen ist.
Deshalb hatten wir für 3 Wochen Trocknungsgeräte im Keller stehen.
Anbei erhalten sie den Energieverbrauch/-nachweis der Geräte mit der Bitte, mir für diesen eine Zwischenrechnung zukommen zu lassen. Diese kann ich dann bei meiner Versicherung einreichen, welche die Kosten übernimmt. 
Mein Vertragskonto lautet 405604519.
Aufgrund der Einfachheit halber, reicht mir die Rechnung per Email.
Sollten sie noch Fragen haben, können sie sich jederzeit gerne bei mir melden.
Mit freundlichen Grüßen 
Wladimir Klotz

""",
        "reference": "Der Kunde bittet um eine Zwischenrechnung für den Energieverbrauch der Trocknungsgeräte zur Einreichung bei der Versicherung."
    },
    "Sample2": {
        "email": """
Vertragsnummer: 403755518
Canan Schenk
Verbraucherstelle: Hans-Henning-Wilmsen-Weg 958, 74815 Osterburg
Zähler: 92739304
Sehr geehrte Damen und Herren,
Ich kündige den Vertrag wegen dem Umzug zum 30.11.23
Um eine schriftliche Bestätigung wird gebeten.
Neue Adresse:
Dippelgasse 83
74815 Osterburg
Viele Grüße
Canan Schenk

""",
        "reference": "Der Kunde kündigt den Vertrag wegen Umzug und erwartet eine schriftliche Bestätigung dazu."
    }
}

@st.cache_resource
def load_all_models():
    """Pre-load all models during app initialization"""
    with st.spinner("Loading models... This may take a few minutes"):
        models = {
            "mt5-small": pipeline(
                "summarization", 
                model="ak2603/mt5-small-synthetic-data-plus-translated"
            ),
            "Llama 3.2": load_llama_model()
        }
    return models

# Initialize models when app loads
models = load_all_models()

# Streamlit UI Configuration
st.title("📧 Automated Email Summarization")

# Sidebar Controls
with st.sidebar:
    st.header("Configuration")
    model_choice = st.selectbox(
        "Select Model",
        ["mt5-small", "Llama 3.2"],
        index=0
    )
    
    st.markdown("---")
    st.markdown("**Model Information:**")
    st.info(f"Selected model: {model_choice}")
    st.info(f"Total loaded models: {len([m for m in models.values() if m is not None])}")

# Main Content Area
col1, col2 = st.columns([2, 1])

with col1:
    st.subheader("Input Email")
    
    # Example selection
    example_choice = st.selectbox(
        "Choose an example email:",
        ["Custom Email"] + list(EXAMPLES.keys()),
        index=0
    )
    
    # Set email text based on selection
    email_text = EXAMPLES[example_choice]["email"] if example_choice in EXAMPLES else ""
    
    email_input = st.text_area(
        "Paste your email here:",
        value=email_text,
        height=300,
        key="input_text",
        placeholder="Enter email content here..."
    )

with col2:
    st.subheader("Summary Generation")
    if st.button("Generate Summary", use_container_width=True):
        if not email_input:
            st.error("Please enter some email content first!")
        else:
            try:
                selected_model = models[model_choice]

                with st.spinner("Generating summary..."):
                    if model_choice == "mt5-small":
                        result = selected_model(
                            email_input,
                            max_length=150,
                            do_sample=True,
                            repetition_penalty=1.5
                        )[0]['summary_text']
                    elif model_choice == "Llama 3.2":
                        model_obj, tokenizer = selected_model
                        result = generate_llama_summary(
                            email_input,
                            model_obj,
                            tokenizer,
                            PROMPT_TEMPLATE
                        )
                    
                    # Display results
                    st.success("**Generated Summary:**")
                    st.write(result)
                    
                    # Show reference summary if example is selected
                    if example_choice in EXAMPLES:
                        st.success("**Reference Summary:**")
                        st.write(EXAMPLES[example_choice]["reference"])
                    
                    # Add export options
                    st.download_button(
                        label="Download Summary",
                        data=result,
                        file_name="email_summary.txt",
                        mime="text/plain"
                    )
                        
            except Exception as e:
                st.error(f"Error generating summary: {str(e)}")

# Footer
st.markdown("---")
st.markdown("_Automated email summarization system v1.0_")