# app.py
import streamlit as st
from transformers import pipeline

@st.cache_resource
def load_model(model_name):
    if model_name == "mt5-small":
        return pipeline("summarization", model="ak2603/mt5-small-synthetic-data-plus-translated")
    # Add space for other models here
    elif model_name == "Llama 3.2":
        return pipeline("text-generation", model="Walid777/llama3-8b-emails-summarization")              
    elif model_name == "Llama 7b Instruct":
        return None  # Placeholder for future implementation
    else:
        raise ValueError("Model not supported")

# Sample emails with reference summaries (for demonstration only)
SAMPLES = {
    "Sample 1": (
            """
            Sehr geehrte Damen und Herren,
            Ich bitte um die Kündigung meines Vertrages, da ich umziehe.
            Vertragsnummer: 40887935006
            Zählernummer: 17760731625925
            Mit freundlichen Grüßen
            Falk Rosemann, Truppring 7, 02044 Wernigerode
            """,
        "Der Kunde bittet um die Kündigung seines Vertrages aufgrund eines Umzugs und gibt die Vertrags- und Zählernummer an."
    ),
    "Sample 2": (
            """
            Die versprochene Rabattaktion wurde in meiner letzten Rechnung nicht berücksichtigt.
            Mit freundlichen Grüßen
            Prof. Vinko Caspar B.Eng., Jolanda-Pruschke-Platz 835, 42943 Neustadtner Waldnaab
            """,
        "Der Kunde erwähnt, dass die versprochene Rabattaktion in der letzten Rechnung nicht berücksichtigt wurde."
    ),
    "Sample 3": (
            """
            Sehr geehrte Damen und Herren,
            Ich habe vor zwei Wochen eine Rechnung erhalten, die ich nicht nachvollziehen kann. Bitte erklären Sie die Details.
            Herzliche Grüße
            Kirstin Girschner-Meister, Oderwaldallee 510, 35412 Halberstadt
            """,
        "Der Kunde erwähnt, dass die versprochene Rabattaktion in der letzten Rechnung nicht berücksichtigt wurde und erwartet eine Überprüfung und Korrektur der Rechnung."
       
    )
}

# UI Layout
st.title("🇩🇪 German Email Summarizer")
st.markdown("Fine-tuned summarization for German emails")

# Sidebar for model selection and sample emails
with st.sidebar:
    st.header("⚙️ Settings")
    model_choice = st.selectbox("Choose Model", ["mt5-small", "Llama 3.2", "Llama 7b Instruct"], index=0)
    sample_choice = st.selectbox("Try Sample Email", ["Custom Input"] + list(SAMPLES.keys()))

# Load the selected model
summarizer = load_model(model_choice)

# Main interface
col1, col2 = st.columns(2)
with col1:
    if sample_choice == "Custom Input":
        input_text = st.text_area("Input Email", height=300, placeholder="Paste your email here...")
    else:
        input_text = st.text_area("Input Email", value=SAMPLES[sample_choice][0], height=300)

with col2:
    if st.button("Generate Summary"):
        if summarizer is None:
            st.error("Selected model is not implemented yet.")
        else:
            with st.spinner("Generating summary..."):
                try:
                    # Generate summary
                    summary_output = summarizer(
                        input_text,
                        max_length=150,
                        do_sample=True,
                        repetition_penalty=1.5
                    )[0]

                    # Dynamically select key based on pipeline task
                    result_key = 'summary_text' if summarizer.task == 'summarization' else 'generated_text'
                    result = summary_output[result_key]
                    
                    st.success("**Generated Summary:**")
                    st.write(result)
                    
                    # Show sample comparison only if a sample is selected
                    if sample_choice != "Custom Input" and sample_choice in SAMPLES:
                        st.divider()
                        st.markdown(f"**Sample Reference Summary ({sample_choice}):**")
                        st.write(SAMPLES[sample_choice][1])
                        
                except Exception as e:
                    st.error(f"Error generating summary: {str(e)}")