import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("sagorsarker/emailgenerator")
model = AutoModelForCausalLM.from_pretrained("sagorsarker/emailgenerator")

# Streamlit UI styling
st.set_page_config(page_title="Email Generator", layout="centered")

# Add some custom CSS for styling
st.markdown("""
    <style>
        .title {
            text-align: center;
            color: #4B0082;
            font-size: 32px;
            font-weight: bold;
        }
        .header {
            text-align: center;
            color: #6A5ACD;
            font-size: 24px;
        }
        .textarea {
            margin-bottom: 20px;
        }
        .button {
            background-color: #6A5ACD;
            color: white;
            font-weight: bold;
            border-radius: 5px;
            padding: 10px 20px;
        }
        .button:hover {
            background-color: #4B0082;
        }

    </style>
""", unsafe_allow_html=True)

# Title and Header
st.markdown('<div class="title">Email Generator</div>', unsafe_allow_html=True)

# User input for email prompt
user_input = st.text_area("Enter the email content prompt:", height=150, key="email_prompt", max_chars=500)

# Add a styled button
generate_button = st.button("Generate Email", key="generate_button", help="Click to generate email", use_container_width=True)

# Handling the generation
if generate_button:
    if user_input:
        # Tokenize the input
        inputs = tokenizer.encode(user_input, return_tensors="pt")
        
        # Generate the email text
        outputs = model.generate(inputs, max_length=300, num_return_sequences=1, no_repeat_ngram_size=2)
        
        # Decode and display the result
        generated_email = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Display the generated email with some styling
        st.markdown('<div class="header">Generated Email:</div>', unsafe_allow_html=True)
        st.markdown(f'<div>{generated_email}</div>', unsafe_allow_html=True)
    else:
        st.error("Please enter a prompt to generate the email.")