import gradio as gr
import requests
import os

# Define API parameters
API_URL = "https://api-inference.huggingface.co/models/tiiuae/falcon-mamba-7b"
API_KEY = os.getenv("HUGGINGFACE_TOKEN")

# Ensure the token is available
if not API_KEY:
    raise ValueError("Hugging Face API token not found. Please set HUGGINGFACE_TOKEN environment variable.")

# Set up headers for Hugging Face API authentication
headers = {
    "Authorization": f"Bearer {API_KEY}"
}

# Function to query the model
def query_model(user_input):
    payload = {
        "inputs": user_input,
        "parameters": {
            "temperature": 0.7,
            "max_length": 150
        }
    }
    response = requests.post(API_URL, headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()[0]['generated_text']
    else:
        return f"Error {response.status_code}: {response.text}"

# Chatbot function that manages conversation history
def chatbot(input_text, history=[]):
    if input_text.lower() in ["exit", "quit"]:
        return "Take care! Remember, seeking support is a strength.", history

    # Append the user's message to the history
    history.append(("You", input_text))
    
    # Get the model's response
    response = query_model(input_text)
    
    # Append the model's response to the history
    history.append(("Bot", response))
    
    # Return the response and updated history for the UI
    return response, history

# Gradio UI Layout
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🧘‍♀️ Mental Health Chatbot
        ### Hi! I'm here to listen and provide support. How can I help you today?
        """
    )
    
    with gr.Row():
        chatbot_output = gr.Chatbot(label="Chatbot", value=[])
        
    with gr.Row():
        user_input = gr.Textbox(label="Your Message", placeholder="Type your message here...", lines=2)
        send_button = gr.Button("Send")

    # Update chatbot output when the user submits a message
    def respond(user_input, history):
        response, history = chatbot(user_input, history)
        return history, gr.update(value="")

    # Clear the input box after sending the message
    send_button.click(respond, inputs=[user_input, chatbot_output], outputs=[chatbot_output, user_input])

# Launch the Gradio app
demo.launch()