import os
import asyncio
import gradio as gr
import whisper
import edge_tts

from groq import Groq


# =========================
# LOAD WHISPER MODEL
# =========================

model = whisper.load_model("base")


# =========================
# GROQ API SETUP
# =========================

# Replace with your Groq API key
GROQ_API_KEY = "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"

client = Groq(api_key=GROQ_API_KEY)


# =========================
# MAIN VOICE CHAT FUNCTION
# =========================

def voice_chat(audio):

    # Check audio input
    if audio is None:
        return None, "No audio received. Please record again."

    try:

        # =========================
        # STEP 1: SPEECH TO TEXT
        # =========================

        result = model.transcribe(audio)

        user_text = result["text"]

        print("\nUser:", user_text)

        # =========================
        # STEP 2: LLM RESPONSE
        # =========================

        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You are a helpful AI voice assistant."
                },
                {
                    "role": "user",
                    "content": user_text
                }
            ],
            model="llama-3.3-70b-versatile"
        )

        response_text = chat_completion.choices[0].message.content

        print("\nAssistant:", response_text)

        # =========================
        # STEP 3: TEXT TO SPEECH
        # =========================

        output_audio = "response.mp3"

        async def text_to_speech():

            communicate = edge_tts.Communicate(
                text=response_text,
                voice="en-US-AriaNeural"
            )

            await communicate.save(output_audio)

        asyncio.run(text_to_speech())

        # =========================
        # RETURN OUTPUTS
        # =========================

        return output_audio, response_text

    except Exception as e:

        print("Error:", str(e))

        return None, f"Error: {str(e)}"


# =========================
# GRADIO INTERFACE
# =========================

iface = gr.Interface(
    fn=voice_chat,

    inputs=gr.Audio(
        sources=["microphone"],
        type="filepath",
        label="Speak Here"
    ),

    outputs=[
        gr.Audio(label="Assistant Voice"),
        gr.Textbox(label="Assistant Text")
    ],

    title="Realtime Voice Assistant",

    description="""
    Speak with an AI assistant using:
    Whisper + Groq LLM + Edge-TTS
    """,

    flagging_mode="never"
)


# =========================
# LAUNCH APPLICATION
# =========================

if __name__ == "__main__":
    iface.launch()