import os
import gradio as gr
import whisper
from gtts import gTTS
from groq import Groq
from dotenv import load_dotenv
import tempfile

# Load environment variables from .env file
load_dotenv()

# Initialize Whisper model
print("Loading Whisper model...")
whisper_model = whisper.load_model("base")

# Initialize Groq API
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not GROQ_API_KEY:
    raise ValueError("GROQ_API_KEY environment variable not set. Please add it to your environment variables or .env file.")
client = Groq(api_key=GROQ_API_KEY)

# Function to transcribe audio to text
def transcribe_audio(audio_file):
    try:
        result = whisper_model.transcribe(audio_file)
        return result["text"]
    except Exception as e:
        return f"Error in transcription: {e}"

# Function to get response from LLM using Groq API
def get_llm_response(user_input):
    try:
        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
    except Exception as e:
        return f"Error in LLM interaction: {e}"

# Function to convert text to speech
def text_to_speech(text):
    try:
        tts = gTTS(text)
        temp_file = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
        tts.save(temp_file.name)
        return temp_file.name
    except Exception as e:
        return f"Error in text-to-speech conversion: {e}"

# Main chatbot pipeline
def chatbot_pipeline(audio_file):
    # Step 1: Transcribe audio
    user_input = transcribe_audio(audio_file)
    if "Error" in user_input:
        return user_input, "No response", None

    # Step 2: Get response from LLM
    response = get_llm_response(user_input)
    if "Error" in response:
        return user_input, response, None

    # Step 3: Convert LLM response to audio
    response_audio = text_to_speech(response)
    if "Error" in response_audio:
        return user_input, response, None

    return user_input, response, response_audio

# Gradio Interface
interface = gr.Interface(
    fn=chatbot_pipeline,
    inputs=gr.Audio(type="filepath"),
    outputs=[
        gr.Textbox(label="Transcribed Text"),
        gr.Textbox(label="Chatbot Response"),
        gr.Audio(label="Response Audio"),
    ],
    title="Real-Time Voice-to-Voice Chatbot",
    description=(
        "This chatbot transcribes your voice input using Whisper, "
        "processes your input through Groq's API to generate a response, "
        "and converts the response back to speech using GTTS."
    ),
    live=True,
)

if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860)