import os
import streamlit as st
import google.generativeai as genai

# Access the API key as an environment variable from Hugging Face secrets.
# Then, configure the official Gemini client with the API key.
api_key = os.getenv("MentalHealth")
genai.configure(api_key=api_key)



if "messages" not in st.session_state:
    # Initialize the session state for the chat history
    st.session_state.messages = []
    # Gemini models have different roles, so we use 'user' and 'model'.
    # A system message is not directly supported, so we will handle the persona
    # in the prompt or in the response generation logic.

for message in st.session_state.messages:
    # Display existing messages from the session state
    with st.chat_message(message["role"]):
        st.markdown(message["parts"][0])

if prompt := st.chat_input("Type your thoughts here..."):
    # Append the user's message to the chat history and display it
    user_message = {"role": "user", "parts": [prompt]}
    st.session_state.messages.append(user_message)
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            # Prepare the list of messages for the Gemini model.
            # We add a preamble to maintain the therapist persona.
            chat_history_for_gemini = [
                {"role": "user", "parts": ["You are a supportive therapist AI. All your responses should be in this persona."]},
                {"role": "model", "parts": ["Understood. I will respond as a supportive therapist."]}
            ] + st.session_state.messages

            # Initialize the model and generate a response
            model = genai.GenerativeModel('gemini-1.5-flash-latest')
            try:
                # Use the new API syntax to create a completion.
                # The model automatically handles the chat history.
                response = model.generate_content(chat_history_for_gemini, stream=True)
                
                full_reply_content = ""
                # Stream the response to the screen for a better user experience.
                for chunk in response:
                    # Check if the chunk has text before adding it to the reply.
                    if chunk.text:
                        full_reply_content += chunk.text
                        st.markdown(full_reply_content)

            except Exception as e:
                full_reply_content = f"An error occurred: {e}"
                st.markdown(full_reply_content)
                
    # Append the assistant's full response to the session state
    assistant_message = {"role": "model", "parts": [full_reply_content]}
    st.session_state.messages.append(assistant_message)