import gradio as gr
import torch
import whisper
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.agents import initialize_agent, Tool, AgentType
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
from gtts import gTTS
import os
from groq import Groq

# Load Whisper model for transcription
model = whisper.load_model("base")

# Initialize Groq client
client = Groq(api_key="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")

# Function to transcribe audio
def transcribe_audio(audio):
    result = model.transcribe(audio)
    return result["text"]

# Function for text-to-speech conversion
def text_to_speech(text):
    tts = gTTS(text)
    audio_path = "/tmp/response.mp3"
    tts.save(audio_path)
    return audio_path

# Function to interact with Groq API for LLM responses
def get_groq_response(question):
    # Use Groq API to get the answer from LLM
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": question,
            }
        ],
        model="llama-3.3-70b-versatile",
    )
    
    return chat_completion.choices[0].message.content

# Initialize Gradio components
with gr.Blocks() as demo:
    gr.Markdown("# Voice/Text Chatbot with Document-based Q&A and Audio Transcription")
    
    # Audio Input: Microphone for recording or File Upload
    audio_input = gr.Audio(type="filepath", label="Record or Upload Audio")
    
    # Text Input for queries
    text_input = gr.Textbox(label="Enter your question", placeholder="Ask a question based on the document...")
    
    # Output Text
    output_text = gr.Textbox(label="Answer")
    
    # Output Audio (for voice-based response)
    output_audio = gr.Audio(label="Voice Response", type="filepath")
    
    # Button to process the input and generate output
    def process_input(audio_input, text_input):
        if audio_input:
            question = transcribe_audio(audio_input)
        else:
            question = text_input
        
        # Get the answer from the LLM via Groq API
        answer = get_groq_response(question)
        
        # Convert the answer to speech and return both text and audio
        audio_path = text_to_speech(answer)
        return answer, audio_path
    
    # Bind the function to the interface
    audio_input.change(process_input, inputs=[audio_input, text_input], outputs=[output_text, output_audio])
    text_input.submit(process_input, inputs=[audio_input, text_input], outputs=[output_text, output_audio])
    
demo.launch(debug=True)