import os
import gradio as gr
import numpy as np
from gtts import gTTS
import torch
import whisper  # Correct import from openai-whisper package
from groq import Groq
import io
import tempfile  # To handle temporary audio file saving

# Initialize Groq API client
client = Groq(api_key="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")
# Load Whisper model
whisper_model = whisper.load_model("base")  # Use 'whisper' directly

def transcribe_audio(audio_file):
    # Load audio
    audio, sr = sf.read(audio_file)
    # Transcribe audio using Whisper
    result = whisper_model.transcribe(audio, language="en")
    return result['text']

def get_response(prompt):
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama3-8b-8192",
    )
    return chat_completion.choices[0].message.content

def text_to_speech(text):
    tts = gTTS(text)
    audio_buffer = io.BytesIO()
    tts.save(audio_buffer)
    audio_buffer.seek(0)
    return audio_buffer

def chatbot(audio_file):
    # Transcribe audio to text
    user_input = transcribe_audio(audio_file)
    # Get response from Llama 8B
    response = get_response(user_input)
    # Convert response to speech
    audio_output = text_to_speech(response)
    return audio_output

# Gradio interface
iface = gr.Interface(
    fn=chatbot,
    inputs=gr.Audio(type="filepath"),  # Remove 'source' argument
    outputs=gr.Audio(type="filepath"),
    live=True,
    title="Voice to Voice Chatbot",
    description="Speak into the microphone, and the chatbot will respond!"
)

iface.launch()