import gradio as gr
from datasets import load_dataset
import random
import groq  # Assuming you are using the Groq library

# Load dataset
ds = load_dataset("alexandreteles/mental-health-conversational-data")

# Extract columns
context = ds["train"]["Context"]
knowledge = ds["train"]["Knowledge"]
response = ds["train"]["Response"]

# Directly input the Groq API key (replace with your actual API key)
api_key = "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"

# Initialize Groq API client with the API key
client = groq.Client(api_key=api_key)

# Define a function to simulate a conversation
def chatbot(user_input):
    if not user_input.strip():
        return "Please enter a question or concern to receive guidance."

    try:
        # Try to call Groq API to generate a response
        brief_response = client.predict(user_input)  # Ensure that 'predict' or correct method is available
    except Exception as e:
        brief_response = None  # If Groq API fails, set brief_response to None

    if not brief_response:
        # If Groq API does not return a response, fall back to dataset
        idx = random.randint(0, len(context)-1)
        reply = response[idx]
        context_text = context[idx]
        knowledge_text = knowledge[idx]
        
        # Combine dataset info with fallback response
        complete_response = (
            f"**Contextual Information**\n{context_text}\n\n"
            f"**Knowledge Base**\n{knowledge_text}\n\n"
            f"**Fallback Response**\n{reply}"
        )
    else:
        # If Groq API returns a response, use that as the final output
        complete_response = f"**Personalized Response**\n{brief_response}"

    return complete_response

# Enhance the interface
interface = gr.Interface(
    fn=chatbot,
    inputs=gr.Textbox(label="Ask your question:", placeholder="How are you feeling today?"),
    outputs=gr.Markdown(label="Psychologist Assistant Response"),
    title="Virtual Psychologist Assistant",
    description=(
        "This is a supportive assistant designed to provide compassionate guidance "
        "for mental well-being. Type your thoughts or questions for tailored advice and insights."
    ),
    theme="huggingface",  # Optional: apply a theme if available
)

# Launch the app
interface.launch()