import gradio as gr
from groq import Groq
from dotenv import load_dotenv
import os

# Load environment variables from .env file
load_dotenv()

# Access the API key
api_key = os.getenv("MY_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):
    try:
        # Prepare messages with chat history
        messages = [{"role": "system", "content": "You are a helpful mental health assistant."}]
        for user_message, bot_response in chat_history:
            messages.append({"role": "user", "content": user_message})
            messages.append({"role": "assistant", "content": bot_response})
        messages.append({"role": "user", "content": user_input})

        # 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

    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:
        # Variable to store chat history
        chat_history = []

        # Modify `generate_and_clear` to accept additional unused arguments
        def generate_and_clear(user_input, *args, **kwargs):
            response, updated_history = generate_response(user_input, chat_history)
            chat_history.append((user_input, response))
            return response

        def clear_chat():
            chat_history.clear()
            return "Chat history cleared. Start a new conversation!"

        # Title and instructions
        gr.Markdown("## Mental Health Chatbot - Powered by LLaMA on Groq")

        # Chat interface setup
        chatbot = gr.ChatInterface(
            fn=generate_and_clear,
            additional_inputs=[gr.Textbox(placeholder="Enter your message here...", label="Your Message")],
            title="Mental Health Chatbot"
        )

        # Add button to clear chat history
        clear_button = gr.Button("Clear Chat History")
        clear_button.click(fn=clear_chat, inputs=[], outputs=chatbot)  # Clear chat on button click

    demo.launch()

# Run the interface
gradio_interface()