import streamlit as st
import openai
import pandas as pd
import plotly.express as px
from datetime import datetime
import os
import random

# --- GPT-4 API Configuration ---
openai.api_key = os.getenv("Open_AI_GPT_4o")

# --- Page Configuration ---
st.set_page_config(page_title="Emma - Mental Health Companion", page_icon="🌟", layout="wide")

# --- Session State Initialization ---
if 'user_profile' not in st.session_state:
    st.session_state.user_profile = None
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'current_page' not in st.session_state:
    st.session_state.current_page = 'home'

# --- Helper Functions ---
def get_gpt4_response(prompt):
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "You are Emma, a helpful assistant."},
                {"role": "user", "content": prompt}
            ]
        )
        return response["choices"][0]["message"]["content"]
    except Exception as e:
        st.error("Error with GPT-4 API. Please check your API key or connection.")
        return "I'm having trouble responding at the moment. Please try again later."

# --- Onboarding and User Profile ---
def save_user_profile(profile_data):
    """Save user profile to session state"""
    st.session_state.user_profile = profile_data

def onboarding_page():
    if st.session_state.user_profile is None:
        st.title("Welcome to Emma 🌟")
        st.write("Your personal mental health companion")

        with st.form("onboarding_form"):
            name = st.text_input("What should I call you?")
            challenges = st.multiselect(
                "What challenges would you like support with?",
                ["Anxiety", "Depression", "Stress", "Sleep Issues", "Relationship Issues", "Work-Life Balance"]
            )
            approach_style = st.slider("Preferred approach style (0: Reassuring, 10: Solution-Oriented):", 0, 10, 5)
            tone_style = st.slider("Preferred tone (0: Lighthearted, 10: Serious):", 0, 10, 5)

            submitted = st.form_submit_button("Start My Journey")
            if submitted:
                profile_data = {
                    "name": name,
                    "challenges": challenges,
                    "preferences": {
                        "approach_style": approach_style,
                        "tone_style": tone_style
                    },
                    "created_at": str(datetime.now())
                }
                save_user_profile(profile_data)
                st.success("Profile created successfully! Let's begin your journey.")
                st.rerun()
    else:
        st.title(f"Welcome back, {st.session_state.user_profile['name']}! 🌟")
        if st.button("Continue to Chat"):
            st.session_state.current_page = "chat"

# --- Chat Interface ---
def chat_page():
    st.title("Chat with Emma 💭")

    for message in st.session_state.chat_history:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    prompt = st.chat_input("Type your message here...")
    if prompt:
        st.session_state.chat_history.append({"role": "user", "content": prompt})
        emma_response = get_gpt4_response(prompt)
        st.session_state.chat_history.append({"role": "assistant", "content": emma_response})
        st.session_state["needs_rerun"] = True  # Set flag to trigger rerun

    # Check for rerun
    if st.session_state.get("needs_rerun", False):
        st.session_state["needs_rerun"] = False
        st.set_query_params()  # Use st.set_query_params to trigger a rerun



# --- Progress Dashboard ---
def progress_dashboard():
    st.title("Your Progress Dashboard 📊")
    
    # Sample mood data
    dates = pd.date_range(start="2024-01-01", end="2024-01-07", freq="D")
    moods = [random.randint(1, 5) for _ in range(len(dates))]
    
    df = pd.DataFrame({
        "Date": dates,
        "Mood": moods
    })

    fig = px.line(df, x="Date", y="Mood", title="Your Mood Timeline")
    st.plotly_chart(fig)

# --- Support Tools ---
def support_tools():
    st.title("Support Tools & Resources 🛠️")
    st.subheader("Guided Meditation")
    st.write("Select a meditation type to begin.")
    meditation_type = st.selectbox("Meditation Type", ["Stress Relief", "Mindfulness", "Sleep"])
    if st.button("Play Meditation"):
        st.write(f"Playing {meditation_type} meditation 🎵")

# --- Settings Page ---
def settings_page():
    st.title("Settings ⚙️")
    if st.session_state.user_profile:
        st.subheader("Therapist Preferences")
        new_approach = st.slider("Approach Style", 0, 10, st.session_state.user_profile["preferences"]["approach_style"])
        new_tone = st.slider("Tone Style", 0, 10, st.session_state.user_profile["preferences"]["tone_style"])
        if st.button("Update Preferences"):
            st.session_state.user_profile["preferences"].update({
                "approach_style": new_approach,
                "tone_style": new_tone
            })
            st.success("Preferences updated!")

# --- Page Navigation ---
def navigation():
    with st.sidebar:
        st.title("Emma 🌟")
        if st.button("Home"):
            st.session_state.current_page = 'home'
        if st.button("Chat"):
            st.session_state.current_page = 'chat'
        if st.button("Progress Dashboard"):
            st.session_state.current_page = 'progress'
        if st.button("Support Tools"):
            st.session_state.current_page = 'tools'
        if st.button("Settings"):
            st.session_state.current_page = 'settings'

# --- Main App Logic ---
def main():
    navigation()

    # Load pages based on navigation
    if st.session_state.current_page == "home":
        onboarding_page()
    elif st.session_state.current_page == "chat":
        chat_page()
    elif st.session_state.current_page == "progress":
        progress_dashboard()
    elif st.session_state.current_page == "tools":
        support_tools()
    elif st.session_state.current_page == "settings":
        settings_page()

if __name__ == "__main__":
    main()