import os
import streamlit as st
import torch
from transformers import pipeline
import speech_recognition as sr
from gtts import gTTS
from io import BytesIO
from pydub import AudioSegment
import tempfile

# Set Hugging Face API key (replace 'your_hugging_face_api_key' with your actual API key)
os.environ["HUGGINGFACEHUB_API_TOKEN"] = "doctor app"

# Load the Hugging Face model for text generation
try:
    chatbot = pipeline("text-generation", model="thrishala/mental_health_chatbot")
except Exception as e:
    st.error(f"Error loading model: {e}")

# Function to capture voice input using SpeechRecognition and OpenAI Whisper
def get_voice_input():
    recognizer = sr.Recognizer()
    with sr.Microphone() as source:
        st.write("Listening...")
        audio = recognizer.listen(source)
        st.write("Recognizing...")
        try:
            # Converting audio to text using OpenAI Whisper if available, fallback to Google's API otherwise
            text = recognizer.recognize_google(audio)
            return text
        except sr.UnknownValueError:
            st.error("Sorry, I could not understand the audio.")
            return None
        except sr.RequestError as e:
            st.error(f"Could not request results; {e}")
            return None

# Function to generate voice response using gTTS
def speak(text):
    try:
        tts = gTTS(text=text, lang='en')
        audio_file = BytesIO()
        tts.write_to_fp(audio_file)
        audio_file.seek(0)
        return audio_file
    except Exception as e:
        st.error(f"Error generating audio response: {e}")
        return None

# Streamlit app layout
st.title("Mental Health Chatbot")
st.write("Talk to your mental health assistant!")

# Voice input button
if st.button("Speak"):
    user_input = get_voice_input()
    if user_input:
        st.write(f"You: {user_input}")
        # Get response from the chatbot
        try:
            response = chatbot(user_input, max_length=150, num_return_sequences=1)[0]['generated_text']
            st.write(f"Bot: {response}")
            # Generate voice response
            audio_output = speak(response)
            if audio_output:
                st.audio(audio_output, format="audio/mp3")
        except Exception as e:
            st.error(f"Error generating response: {e}")

# Text input
user_input = st.text_input("Type your message:")
if st.button("Send"):
    if user_input:
        st.write(f"You: {user_input}")
        # Get response from the chatbot
        try:
            response = chatbot(user_input, max_length=150, num_return_sequences=1)[0]['generated_text']
            st.write(f"Bot: {response}")
            # Generate voice response
            audio_output = speak(response)
            if audio_output:
                st.audio(audio_output, format="audio/mp3")
        except Exception as e:
            st.error(f"Error generating response: {e}")