import gradio as gr
from groq import Groq
from langchain_google_genai import ChatGoogleGenerativeAI
import os
import tempfile

# -------------------- API Configuration --------------------
# Initialize clients
client = Groq(api_key="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")
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    google_api_key="AIzaSyD2DMFgcL0kWTQYhii8wseSHY3BRGWSebk",
    max_output_tokens=500
)

# -------------------- Core Functions --------------------
def transcribe_audio(audio_path, language="ar"):
    """Transcribe audio file using Groq Whisper"""
    try:
        with open(audio_path, "rb") as audio_file:
            transcription = client.audio.transcriptions.create(
                file=(os.path.basename(audio_path), audio_file.read()),
                model="whisper-large-v3-turbo",
                response_format="verbose_json",
                language=language
            )
        return transcription.text, transcription.language
    except Exception as e:
        return f"Error in transcription: {str(e)}", None

def get_ai_response(text, detected_language="ar"):
    """Get AI response from Gemini with Arabic system prompt"""
    try:
        system_prompt = """أنت مساعد ذكي ومفيد. يجب عليك الرد باللغة العربية واستخدام نفس اللهجة التي يستخدمها المستخدم.
        
قواعد مهمة:
- إذا تحدث المستخدم بالعربية الفصحى، رد بالفصحى
- إذا تحدث المستخدم بلهجة عامية (مصرية، خليجية، شامية، مغاربية، إلخ)، رد بنفس اللهجة
- حافظ على نفس الأسلوب والنبرة التي يستخدمها المستخدم
- كن طبيعياً ومحادثاً ومفيداً
- الرد يجب أن يكون مختصراً ومفيداً وواضحاً."""

        from langchain_core.messages import HumanMessage, SystemMessage
        
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=text)
        ]
        
        response = llm.invoke(messages)
        return response.content
    except Exception as e:
        return f"Error getting AI response: {str(e)}"

def text_to_speech(text, voice_choice="Amira-PlayAI", language="ar"):
    """Convert text to speech using Groq TTS"""
    try:
        # Use selected Arabic voice
        voice = voice_choice
        model = "playai-tts-arabic"
        
        response = client.audio.speech.create(
            model=model,
            voice=voice,
            response_format="mp3",
            input=text,
        )
        
        # Save to temporary file
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        response.write_to_file(temp_file.name)
        return temp_file.name
    except Exception as e:
        print(f"TTS Error: {str(e)}")  # Debug print
        # Return error message instead of None
        return f"Error generating audio: {str(e)}"

# -------------------- Gradio Interface Function --------------------
def process_voice_chat(audio, voice_choice):
    """Main function to process voice input and generate response"""
    
    if audio is None:
        return "يرجى تقديم مدخل صوتي", "", None
    
    # Always use Arabic
    lang_code = "ar"
    
    # Step 1: Transcribe audio
    transcription, detected_lang = transcribe_audio(audio, lang_code)
    
    if transcription.startswith("Error"):
        return transcription, "", None
    
    # Step 2: Get AI response
    ai_response = get_ai_response(transcription, detected_lang or lang_code)
    
    if ai_response.startswith("Error"):
        return transcription, ai_response, None
    
    # Step 3: Convert response to speech with selected voice
    audio_output = text_to_speech(ai_response, voice_choice, "ar")
    
    # Check if audio generation failed
    if isinstance(audio_output, str) and audio_output.startswith("Error"):
        return transcription, ai_response, None
    
    return transcription, ai_response, audio_output

# -------------------- Gradio Interface --------------------
with gr.Blocks(title="مساعد المحادثة الصوتية", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎤 مساعد المحادثة الصوتية
        سجل صوتك أو ارفع ملف صوتي للمحادثة مع الذكاء الاصطناعي
        """
    )
    
    with gr.Row():
        with gr.Column():
            voice_selector = gr.Radio(
                choices=["Amira-PlayAI", "Ahmad-PlayAI"],
                value="Amira-PlayAI",
                label="اختر الصوت | Select Voice",
                info="أميرة (أنثى) | أحمد (ذكر)"
            )
            audio_input = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="سجل أو ارفع ملف صوتي"
            )
            submit_btn = gr.Button("معالجة", variant="primary", size="lg")
        
        with gr.Column():
            transcription_output = gr.Textbox(
                label="رسالتك (النص المنسوخ)",
                lines=3,
                rtl=True
            )
            ai_response_output = gr.Textbox(
                label="رد الذكاء الاصطناعي",
                lines=5,
                rtl=True
            )
            audio_output = gr.Audio(
                label="الرد الصوتي",
                type="filepath",
                autoplay=True
            )
    
    # Button action
    submit_btn.click(
        fn=process_voice_chat,
        inputs=[audio_input, voice_selector],
        outputs=[transcription_output, ai_response_output, audio_output]
    )
    
    gr.Markdown(
        """
        ### 📝 التعليمات:
        1. اختر الصوت المفضل (أميرة أو أحمد)
        2. سجل صوتك باستخدام الميكروفون أو ارفع ملف صوتي
        3. اضغط على "معالجة" للحصول على رد الذكاء الاصطناعي مع صوت
        
        ### 🔑 للنشر على Hugging Face Spaces:
        أضف هذه الأسرار في إعدادات المساحة:
        - `GROQ_API_KEY`: مفتاح Groq API الخاص بك
        - `GOOGLE_API_KEY`: مفتاح Google API الخاص بك
        """
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()