import streamlit as st
import os
from groq import Groq
from streamlit_extras.switch_page_button import switch_page

# Initialize Groq client with API key
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

# Configure Streamlit page settings
st.set_page_config(
    page_title="EmailGenie",
    page_icon="📧",
    layout="centered",
    initial_sidebar_state="collapsed",
)

# Hide the sidebar
st.markdown(
    """
<style>
    [data-testid="collapsedControl"] {
        display: none
    }
</style>
""",
    unsafe_allow_html=True,
)

# Hide Streamlit's default styling
hide_st_style = """
    <style>
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    header {visibility: hidden;}
    </style>"""
st.markdown(hide_st_style, unsafe_allow_html=True)

# Display the page title and description
st.markdown(
    "<h1 style='text-align: center;'>Welcome to 📧EmailGenie</h1>", unsafe_allow_html=True
)
st.markdown(
    "<p style='text-align: center;'>AI-Powered Cold Email Generator</p>",
    unsafe_allow_html=True,
)

# Initialize conversation history in session state if it doesn't exist
if "conversation_history" not in st.session_state:
    st.session_state.conversation_history = []

# Initialize session state variables for input fields
if 'recipient_name' not in st.session_state:
    st.session_state.recipient_name = ""
if 'recipient_role' not in st.session_state:
    st.session_state.recipient_role = ""
if 'industry' not in st.session_state:
    st.session_state.industry = ""
if 'company_name' not in st.session_state:
    st.session_state.company_name = ""
if 'sender_name' not in st.session_state:
    st.session_state.sender_name = ""
if 'sender_role' not in st.session_state:
    st.session_state.sender_role = ""
if 'specific_goal' not in st.session_state:
    st.session_state.specific_goal = ""

# Function to generate the email prompt
def generate_prompt(email_type, recipient_name, industry, recipient_role, company_name, sender_name, sender_role, specific_goal, email_tone, email_style):
    prompt = f"""Instructions to draft the email:
    1. Create a {email_type} email based on the following details.
    2. Generate an appropriate subject line for the email.
    3. Format the email properly with greetings, body, and closing.
    4. Ensure the email has a {email_tone} tone and follows a {email_style} style.

    Email details:
    - Recipient Name: {recipient_name}
    - Recipient Role: {recipient_role}
    - Recipient Company: {company_name}
    - Recipient Industry: {industry}
    - Sender Name: {sender_name}
    - Sender Role: {sender_role}
    - Specific Reason for which sender is reaching out to Recipient: {specific_goal}

    Important:
    - If any of the above fields are empty, do not make assumptions. Use only the provided information.
    - Do not include any commentary or explanations outside of the email content.
    - Avoid using variable names or placeholders in the final email.

    Please generate a personalized {email_type} email using the above details."""

    # Add system prompt to conversation history if it doesn't exist
    if not any(msg['role'] == 'system' for msg in st.session_state.conversation_history):
        st.session_state.conversation_history += [{"role":"system", "content":"You are an expert at writing personalized emails. Don't answer questions that are not related to generating or modifying emails. For irrelevant questions just respond something like you can only help with writing personalized emails."}]

    st.session_state.conversation_history += [{"role": "user", "content": prompt}]

    return prompt

# Function to get response from the AI model
def get_response(prompt):
    messages = st.session_state.conversation_history
    model = 'llama-3.1-8b-instant'

    chat_completion = client.chat.completions.create(
        messages=messages,
        model= model
    )
    
    response = chat_completion.choices[0].message.content
    st.session_state.conversation_history += [{"role": "assistant", "content": response}]
    
    return response

# CSS to make radio button options appear side by side
st.markdown(
"""
<style>
.stRadio > div {
flex-direction: row;
}
</style>
""",
    unsafe_allow_html=True,
)

# Input Fields
st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Email Type</span>", unsafe_allow_html=True)
email_type = st.radio("", options=["Sales Outreach", "Networking", "Job Application", "Partnership Proposal", "Follow-Up"], index=["Sales Outreach", "Networking", "Job Application", "Partnership Proposal", "Follow-Up"].index(st.session_state.get('email_type', "Sales Outreach")), label_visibility="collapsed")

# Create text input fields for user information
st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Recipient Name</span>", unsafe_allow_html=True)
recipient_name = st.text_input("Recipient Name", value=st.session_state.get('recipient_name', ""), placeholder="e.g. John Doe", label_visibility="collapsed")

st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Recipient's Role</span>", unsafe_allow_html=True)
recipient_role = st.text_input("Recipient's Role", value=st.session_state.get('recipient_role', ""), placeholder="e.g. Marketing Manager, CEO, Software Engineer", label_visibility="collapsed")

st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Industry</span>", unsafe_allow_html=True)
industry = st.text_input("Industry", value=st.session_state.get('industry', ""), placeholder="e.g. Technology, Healthcare, Finance", label_visibility="collapsed")

st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Company Name</span>", unsafe_allow_html=True)
company_name = st.text_input("Company Name", value=st.session_state.get('company_name', ""), placeholder="e.g. Acme Corporation", label_visibility="collapsed")

st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Your Name</span>", unsafe_allow_html=True)
sender_name = st.text_input("Your Name", value=st.session_state.get('sender_name', ""), placeholder="e.g. Jane Smith", label_visibility="collapsed")

st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Your Role</span>", unsafe_allow_html=True)
sender_role = st.text_input("Your Role", value=st.session_state.get('sender_role', ""), placeholder="e.g. Sales Representative", label_visibility="collapsed")

st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Specific Goal of the Email</span>", unsafe_allow_html=True)
specific_goal = st.text_input("Specific Goal of the Email", value=st.session_state.get('specific_goal', ""), placeholder="e.g. Schedule a product demo", label_visibility="collapsed")

# Email Tone and Style Selection
st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Email Tone</span>", unsafe_allow_html=True)
email_tone = st.radio("", options=["Professional", "Friendly", "Formal", "Casual", "Enthusiastic"], index=["Professional", "Friendly", "Formal", "Casual", "Enthusiastic"].index(st.session_state.get('email_tone', "Professional")), label_visibility="collapsed")

st.markdown("<span style='color: #c94d4d; font-weight: bold;'>Email Style</span>", unsafe_allow_html=True)
email_style = st.radio("", options=["Concise", "Detailed", "Persuasive", "Informative"], index=["Concise", "Detailed", "Persuasive", "Informative"].index(st.session_state.get('email_style', "Concise")), label_visibility="collapsed")

# Update session state with current input values
st.session_state.email_type = email_type
st.session_state.recipient_name = recipient_name
st.session_state.recipient_role = recipient_role
st.session_state.industry = industry
st.session_state.company_name = company_name
st.session_state.sender_name = sender_name
st.session_state.sender_role = sender_role
st.session_state.specific_goal = specific_goal
st.session_state.email_tone = email_tone
st.session_state.email_style = email_style

# Generate Email Button
if st.button("Generate Email", use_container_width=True, type='primary'):
    # Check if all fields are filled
    required_fields = [
        (st.session_state.email_type, "Email Type"),
        (st.session_state.recipient_name, "Recipient Name"),
        (st.session_state.recipient_role, "Recipient's Role"),
        (st.session_state.industry, "Industry"),
        (st.session_state.company_name, "Company Name"),
        (st.session_state.sender_name, "Your Name"),
        (st.session_state.sender_role, "Your Role"),
        (st.session_state.specific_goal, "Specific Goal of the Email"),
        (st.session_state.email_tone, "Email Tone"),
        (st.session_state.email_style, "Email Style")
    ]
    
    empty_fields = [field_name for field_value, field_name in required_fields if not field_value.strip()]
    
    if empty_fields:
        st.error(f"Please fill in all required fields: {', '.join(empty_fields)}")
    else:
        # Generate email prompt and get AI response
        email_prompt = generate_prompt(st.session_state.email_type, st.session_state.recipient_name, st.session_state.industry, st.session_state.recipient_role, st.session_state.company_name, st.session_state.sender_name, st.session_state.sender_role, st.session_state.specific_goal, st.session_state.email_tone, st.session_state.email_style)
        email_content = get_response(email_prompt)
        
        # Store the generated email in session state
        st.session_state.generated_email = email_content
        
        # Switch to the follow-up page
        switch_page("follow_up")

# Preview generated email if it exists
if 'generated_email' in st.session_state:
    st.subheader("Generated Email Preview")
    st.write(st.session_state.generated_email)