import streamlit as st
from groq import Groq

# Groq API Key (replace with your actual API key)
groq_api_key = "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"

# Initialize Groq client
client = Groq(api_key=groq_api_key)

# Function to fetch personalized advice using Groq API
def get_personalized_advice(user_data):
    """Fetch personalized advice using the Groq API based on user data."""
    query = f"Based on the user's data: anxiety level: {user_data['anxiety_level']}, self esteem: {user_data['self_esteem']}, mental health history: {user_data['mental_health_history']}, stress level: {user_data['stress_level']}, provide personalized mental health advice."
    
    try:
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": query}],
            model="llama-3.3-70b-versatile",
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"Error fetching advice from Groq: {str(e)}")
        return None

# Function for chatbot interaction with professional styling
def chatbot_interface():
    st.set_page_config(page_title="🎓 Student Well-being Chatbot", layout="centered")
    st.title("🎓 Student Well-being Chatbot")
    st.write("Welcome! I'm here to help you with well-being advice tailored to your current state.")

    # Styling the sidebar for improved UX
    st.sidebar.title("Student Well-being")
    st.sidebar.write("The well-being questionnaire will guide you through assessing your mental health and providing personalized advice.")

    # Initialize session state
    if 'step' not in st.session_state:
        st.session_state.step = 0
        st.session_state.user_data = {}

    def next_step():
        st.session_state.step += 1

    def show_input_error(message="Please provide a valid input."):
        st.error(message)
        return False

    # Step 0: Welcome message and initial question
    if st.session_state.step == 0:
        st.write("Bot: Hi there! I'm here to assist you with your well-being. Let's get started! 😊")
        st.markdown("<hr>", unsafe_allow_html=True)
        next_step_button = st.button("Start Chat", key="start_button", help="Click to start the well-being questionnaire.")
        if next_step_button:
            next_step()

    # Step 1: Ask for anxiety level with slider
    elif st.session_state.step == 1:
        st.write("Bot: On a scale of 1 to 10, how would you rate your anxiety level?")
        anxiety_level = st.slider("Your Anxiety Level:", min_value=1, max_value=10, value=5, step=1, help="Rate your anxiety from 1 (low) to 10 (high).")
        st.session_state.user_data['anxiety_level'] = anxiety_level
        next_step()

    # Step 2: Ask for self-esteem with slider
    elif st.session_state.step == 2:
        st.write("Bot: On a scale of 1 to 10, how would you rate your self-esteem?")
        self_esteem = st.slider("Your Self-Esteem Level:", min_value=1, max_value=10, value=5, step=1, help="Rate your self-esteem from 1 (low) to 10 (high).")
        st.session_state.user_data['self_esteem'] = self_esteem
        next_step()

    # Step 3: Ask for mental health history
    elif st.session_state.step == 3:
        st.write("Bot: How would you describe your mental health history?")
        mental_health_history = st.selectbox(
            "Select your mental health history:",
            ["None", "Minor Issues", "Moderate Issues", "Severe Issues"],
            help="Choose an option that best describes your mental health background."
        )
        st.session_state.user_data['mental_health_history'] = mental_health_history
        next_step()

    # Step 4: Ask for stress level with slider
    elif st.session_state.step == 4:
        st.write("Bot: On a scale of 0 to 2, how would you rate your stress level?")
        stress_level = st.slider("Your Stress Level:", min_value=0, max_value=2, value=1, step=1, help="Rate your stress from 0 (low) to 2 (high).")
        st.session_state.user_data['stress_level'] = stress_level
        next_step()

    # Step 5: Show advice based on user input
    elif st.session_state.step == 5:
        st.write("Bot: Thank you for providing the information! Here’s your personalized advice:")
        user_data = st.session_state.user_data
        st.write(f"Your Data: {user_data}")

        # Fetch personalized advice from Groq
        st.subheader("🔔 Here's Your Personalized Well-being Advice:")
        advice = get_personalized_advice(user_data)
        if advice:
            st.write(advice)
        else:
            st.write("I couldn't retrieve specific advice right now, but please check back later.")

        # Option to restart the chat
        st.markdown("<hr>", unsafe_allow_html=True)
        restart_button = st.button("Start Over", key="restart_button", help="Click to restart the questionnaire.")
        if restart_button:
            st.session_state.step = 0
            st.session_state.user_data = {}

# Main function to run the chatbot
def main():
    chatbot_interface()

if __name__ == "__main__":
    main()