import streamlit as st
from huggingface_hub import InferenceClient
from gtts import gTTS
import base64
import os
from langdetect import detect, DetectorFactory

# To ensure consistent results from langdetect
DetectorFactory.seed = 0

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def generate_response(message, system_message, max_tokens, temperature, top_p):
    if not message:
        return "Please enter a message.", None

    messages = [{"role": "system", "content": system_message}]
    messages.append({"role": "user", "content": message})

    response = ""
    for msg in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = msg.choices[0].delta.content
        response += token

    # Detect language and set TTS language
    try:
        lang = detect(message)
    except Exception as e:
        st.error(f"Error detecting language: {e}")
        lang = 'en'

    # Generate speech from the response text using gTTS
    tts = gTTS(text=response, lang='ar' if lang == 'ar' else 'en')
    audio_file = "response.mp3"
    tts.save(audio_file)
    return response, audio_file

def get_base64_audio(audio_file):
    with open(audio_file, "rb") as f:
        audio_data = f.read()
    return base64.b64encode(audio_data).decode()

st.title("A Chatbot with Voice Response")

message = st.text_input("Enter your message:")
system_message = st.text_input("System message:", value="You are a friendly Chatbot.")
max_tokens = st.slider("Max new tokens", 1, 2048, 512)
temperature = st.slider("Temperature", 0.1, 4.0, 0.7, step=0.1)
top_p = st.slider("Top-p (nucleus sampling)", 0.1, 1.0, 0.95, step=0.05)

if st.button("Generate Response"):
    response, audio_file = generate_response(message, system_message, max_tokens, temperature, top_p)
    st.write(response)

    if audio_file:
        # Get base64 encoded audio data
        audio_base64 = get_base64_audio(audio_file)
        
        # HTML to autoplay the audio
        audio_html = f"""
        <audio controls autoplay>
        <source src="data:audio/mp3;base64,{audio_base64}" type="audio/mp3">
        Your browser does not support the audio element.
        </audio>
        """
        st.markdown(audio_html, unsafe_allow_html=True)

        # Clean up the audio file
        os.remove(audio_file)