import os
import streamlit as st
import whisper
from gtts import gTTS
from groq import Groq
from dotenv import load_dotenv
import tempfile
from streamlit_webrtc import WebRtcMode, webrtc_streamer, AudioProcessorBase, ClientSettings
import numpy as np
import wave

# Load environment variables from .env file
load_dotenv()

# Load Whisper model for speech-to-text
st.write("Loading Whisper model...")
whisper_model = whisper.load_model("base")

# Initialize Groq API
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
    raise ValueError("API key is not set. Please ensure GROQ_API_KEY is defined in the .env file.")
st.write("Initializing Groq API...")
client = Groq(api_key=api_key)

# Function to transcribe audio using Whisper
def transcribe_audio(audio_file):
    result = whisper_model.transcribe(audio_file)
    return result['text']

# Function to get LLM response using Groq
def get_llm_response(user_input):
    chat_completion = client.chat.completions.create(
        messages=[
            {"role": "user", "content": user_input}
        ],
        model="llama3-8b-8192",
        stream=False,
    )
    return chat_completion.choices[0].message.content

# Function to convert text response to speech using GTTS
def text_to_speech(text):
    tts = gTTS(text)
    temp_file = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
    tts.save(temp_file.name)
    return temp_file.name

# AudioProcessor class for recording microphone input
class AudioProcessor(AudioProcessorBase):
    def __init__(self):
        self.frames = []

    def recv(self, frame):
        self.frames.append(frame.to_ndarray().flatten())
        return frame

    def save_audio(self, path):
        st.write(f"Saving audio to: {path}")
        with wave.open(path, "wb") as wf:
            wf.setnchannels(1)  # Mono audio
            wf.setsampwidth(2)  # 16-bit audio
            wf.setframerate(16000)  # 16 kHz
            wf.writeframes(np.concatenate(self.frames).tobytes())
        st.write("Audio saved successfully.")

# Streamlit UI
st.title("Real-Time Voice-to-Voice Chatbot")
st.write("Use your microphone or upload an audio file to interact with the chatbot.")

# Microphone option
audio_processor = None
microphone_mode = st.radio("Choose input method:", ["Microphone", "Upload Audio"])

if microphone_mode == "Microphone":
    webrtc_ctx = webrtc_streamer(
        key="speech-to-text",
        mode=WebRtcMode.SENDONLY,
        audio_processor_factory=AudioProcessor,
        client_settings=ClientSettings(
            rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
            media_stream_constraints={"audio": True, "video": False},
        ),
    )

    if webrtc_ctx.audio_processor:
        audio_processor = webrtc_ctx.audio_processor

    if st.button("Process Microphone Input"):
        if audio_processor:
            # Save the microphone input to a temporary file
            audio_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
            try:
                audio_processor.save_audio(audio_path)
                st.write("Microphone input saved.")
            except Exception as e:
                st.error(f"Error saving audio: {e}")
                st.stop()

            # Process the audio
            try:
                with st.spinner("Transcribing audio..."):
                    user_input = transcribe_audio(audio_path)
                    st.write("**Transcribed Text:**")
                    st.write(user_input)
            except Exception as e:
                st.error(f"Error during transcription: {e}")
                st.stop()

            # Generate LLM response
            try:
                with st.spinner("Generating response..."):
                    response = get_llm_response(user_input)
                    st.write("**Chatbot Response:**")
                    st.write(response)
            except Exception as e:
                st.error(f"Error during LLM response generation: {e}")
                st.stop()

            # Convert to audio
            try:
                with st.spinner("Converting response to audio..."):
                    response_audio = text_to_speech(response)
                    st.audio(response_audio, format="audio/mp3")
            except Exception as e:
                st.error(f"Error during text-to-speech conversion: {e}")
                st.stop()
else:
    # File uploader for audio input
    uploaded_audio = st.file_uploader("Upload an audio file (.wav or .mp3)", type=["wav", "mp3"])

    if uploaded_audio is not None:
        st.write("**Processing your input...**")
        try:
            with st.spinner("Transcribing audio..."):
                user_input = transcribe_audio(uploaded_audio)
                st.write("**Transcribed Text:**")
                st.write(user_input)
        except Exception as e:
            st.error(f"Error during transcription: {e}")
            st.stop()

        try:
            with st.spinner("Generating response..."):
                response = get_llm_response(user_input)
                st.write("**Chatbot Response:**")
                st.write(response)
        except Exception as e:
            st.error(f"Error during LLM response generation: {e}")
            st.stop()

        try:
            with st.spinner("Converting response to audio..."):
                response_audio = text_to_speech(response)
                st.audio(response_audio, format="audio/mp3")
        except Exception as e:
            st.error(f"Error during text-to-speech conversion: {e}")
            st.stop()

st.write("---")
st.write("Developed using Whisper, Groq, and GTTS")