import gradio as gr
from groq import Groq
import os
import time

api_key = os.getenv('GROQ_API_KEY')
# Initialize Groq client
client = Groq(api_key=api_key)

# Function to generate responses with error handling
def generate_response(user_input, chat_history: list):
    try:
        # Prepare messages with chat history
        messages = [{"role": "system", "content": "You are a mental health assistant. Your responses should be empathetic, non-judgmental, and provide helpful advice based on mental health principles. Always encourage seeking professional help when needed. Your responses should look human as well."}]
        
        # Iterate through chat history and add user and assistant messages
        for message in chat_history:
            # Ensure that each message contains only 'role' and 'content' keys
            if 'role' in message and 'content' in message:
                messages.append({"role": message["role"], "content": message["content"]})
            else:
                print(f"Skipping invalid message: {message}")
        
        messages.append({"role": "user", "content": user_input})  # Add the current user message

        # Call Groq API to get a response from LLaMA
        chat_completion = client.chat.completions.create(
            messages=messages,
            model="llama3-8b-8192"
        )
        
        # Extract response
        response = chat_completion.choices[0].message.content
        return response, chat_history  # Ensure you return both response and chat_history

    except Exception as e:
        print(f"Error occurred: {e}")  # Print error to console for debugging
        return "An error occurred while generating the response. Please try again.", chat_history

# Define Gradio interface
def gradio_interface():
    with gr.Blocks() as demo:
        # Initialize chat history
        chat_history = []

        # Create input textbox and button for clearing chat
        gr.Markdown("## Mental Health Chatbot")
        chatbot = gr.Chatbot(type="messages")
        msg = gr.Textbox(placeholder="Type your message here...")
        clear = gr.Button("Clear")

        # User message submission function
        def user(user_message, history: list):
            # Add user message to the history
            history.append({"role": "user", "content": user_message})
            return "", history  # Reset message input and return updated history

        # Bot response function with simulated typing effect
        def bot(history: list):
            # Ensure that history is not empty
            if len(history) > 0:
                user_input = history[-1]["content"]  # Get the last user message
                response, updated_history = generate_response(user_input, history)  # Get bot's response

                history = updated_history  # Update the history with the new response
                history.append({"role": "assistant", "content": ""})  # Add placeholder for assistant

                # Simulate typing effect for the bot's response
                for character in response:
                    history[-1]['content'] += character
                    time.sleep(0.02)  # Typing delay
                    yield history  # Yield updated history as the bot types

        # Set up interaction flow:
        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(bot, chatbot, chatbot)
        clear.click(lambda: [], None, chatbot, queue=False)  # Clear chat history when clicked

    demo.launch()

# Run the interface
gradio_interface()